Repository: ggerganov/llama.cpp
File: ggml/src/ggml-cuda/ggml-cuda.cu (CUDA)

Code before:
    // Only norm/neox shaders have the fusion code
    const int mode = ((const int32_t *) rope->op_params)[2];
    if (mode != GGML_ROPE_TYPE_NORMAL && mode != GGML_ROPE_TYPE_NEOX) {
        return false;
    }

    return true;
}

static bool ggml_cuda_can_fuse(const struct ggml_cgraph * cgraph, int node_idx, std::initializer_list<enum ggml_op> ops, std::initializer_list<enum ggml_unary_op> unary_ops) {
#ifndef NDEBUG
    const size_t num_unary = std::count(ops.begin(), ops.end(), GGML_OP_UNARY);
    GGML_ASSERT(unary_ops.size() == num_unary);
#endif

    //TODO: remove special case once ggml_can_fuse can handle empty nodes
    std::initializer_list<enum ggml_op> topk_moe_ops =
        ggml_cuda_topk_moe_ops(/*with_norm*/ false, /*delayed_softmax=*/false);
    std::initializer_list<enum ggml_op> topk_moe_ops_with_norm =
        ggml_cuda_topk_moe_ops(/*with_norm=*/true, /*delayed_softmax=*/false);
    std::initializer_list<enum ggml_op> topk_moe_ops_delayed_softmax =
        ggml_cuda_topk_moe_ops(/*with_norm=*/false, /*delayed_softmax=*/true);

    const auto is_equal = [](const std::initializer_list<enum ggml_op> & list1,
                             const std::initializer_list<enum ggml_op> & list2) {
        return list1.size() == list2.size() && std::equal(list1.begin(), list1.end(), list2.begin(), list2.end());
    };

    if (is_equal(topk_moe_ops_with_norm, ops) &&
        ggml_can_fuse_subgraph(cgraph, node_idx, ops, { node_idx + 3, node_idx + 9 })) {
        ggml_tensor * softmax = cgraph->nodes[node_idx];
        ggml_tensor * weights = cgraph->nodes[node_idx + 9];

        if (ggml_cuda_should_use_topk_moe(softmax, weights)) {
            return true;
        }
    }

    if (is_equal(topk_moe_ops, ops) && ggml_can_fuse_subgraph(cgraph, node_idx, ops, { node_idx + 3, node_idx + 4 })) {
        ggml_tensor * softmax = cgraph->nodes[node_idx];
        ggml_tensor * weights = cgraph->nodes[node_idx + 4];
        if (ggml_cuda_should_use_topk_moe(softmax, weights)) {
            return true;
        }
    }

    if (is_equal(topk_moe_ops_delayed_softmax, ops) &&
        ggml_can_fuse_subgraph(cgraph, node_idx, ops, { node_idx + 1, node_idx + 5 })) {
        ggml_tensor * softmax = cgraph->nodes[node_idx + 4];
        ggml_tensor * weights = cgraph->nodes[node_idx + 5];


Review comment (suggestion):
```suggestion
        return list1.size() == list2.size() && std::equal(list1.begin(), list1.end(), list2.begin(), list2.end());
```
From implementation specs, I think size comparison is redundant when you are passing end for list2, I don't know if there's any difference between checking size and not passing end or not checking size and passing end though. :)

Code after:
    // Only norm/neox shaders have the fusion code
    const int mode = ((const int32_t *) rope->op_params)[2];
    if (mode != GGML_ROPE_TYPE_NORMAL && mode != GGML_ROPE_TYPE_NEOX) {
        return false;
    }

    return true;
}

static bool ggml_cuda_can_fuse(const struct ggml_cgraph * cgraph, int node_idx, std::initializer_list<enum ggml_op> ops, std::initializer_list<enum ggml_unary_op> unary_ops) {
#ifndef NDEBUG
    const size_t num_unary = std::count(ops.begin(), ops.end(), GGML_OP_UNARY);
    GGML_ASSERT(unary_ops.size() == num_unary);
#endif

    //TODO: remove special case once ggml_can_fuse can handle empty nodes
    std::initializer_list<enum ggml_op> topk_moe_ops =
        ggml_cuda_topk_moe_ops(/*with_norm*/ false, /*delayed_softmax=*/false);
    std::initializer_list<enum ggml_op> topk_moe_ops_with_norm =
        ggml_cuda_topk_moe_ops(/*with_norm=*/true, /*delayed_softmax=*/false);
    std::initializer_list<enum ggml_op> topk_moe_ops_delayed_softmax =
        ggml_cuda_topk_moe_ops(/*with_norm=*/false, /*delayed_softmax=*/true);

    const auto is_equal = [](const std::initializer_list<enum ggml_op> & list1,
                             const std::initializer_list<enum ggml_op> & list2) {
        return std::equal(list1.begin(), list1.end(), list2.begin(), list2.end());
    };

    if (is_equal(topk_moe_ops_with_norm, ops) &&
        ggml_can_fuse_subgraph(cgraph, node_idx, ops, { node_idx + 3, node_idx + 9 })) {
        ggml_tensor * softmax = cgraph->nodes[node_idx];
        ggml_tensor * weights = cgraph->nodes[node_idx + 9];

        if (ggml_cuda_should_use_topk_moe(softmax, weights)) {
            return true;
        }
    }

    if (is_equal(topk_moe_ops, ops) && ggml_can_fuse_subgraph(cgraph, node_idx, ops, { node_idx + 3, node_idx + 4 })) {
        ggml_tensor * softmax = cgraph->nodes[node_idx];
        ggml_tensor * weights = cgraph->nodes[node_idx + 4];
        if (ggml_cuda_should_use_topk_moe(softmax, weights)) {
            return true;
        }
    }

    if (is_equal(topk_moe_ops_delayed_softmax, ops) &&
        ggml_can_fuse_subgraph(cgraph, node_idx, ops, { node_idx + 1, node_idx + 5 })) {
        ggml_tensor * softmax = cgraph->nodes[node_idx + 4];
        ggml_tensor * weights = cgraph->nodes[node_idx + 5];