Function compute_moe_layer_node

cactus/graph/graph_ops_nn.cpp:340–512

}

void compute_moe_layer_node(GraphNode& node, const std::vector<std::unique_ptr<GraphNode>>& nodes, const std::unordered_map<size_t, size_t>& node_index_map) {
    const size_t num_experts = node.params.num_experts;
    const size_t top_k = node.params.num_experts_per_tok;
    const bool normalize_routing = node.params.normalize_routing;
    const float eps = node.params.epsilon;
    const float routed_scaling_factor = node.params.scalar;
    const bool gated = node.params.moe_gated;
    const Activation activation = node.params.activation;
    const size_t base_inputs = gated ? (3 + 3 * num_experts) : (3 + 2 * num_experts);
    bool has_per_expert_scale = node.input_ids.size() == base_inputs + 1;
    if (node.input_ids.size() != base_inputs && node.input_ids.size() != base_inputs + 1) {
        throw std::runtime_error("moe_layer expects " + std::to_string(base_inputs) + " or " + std::to_string(base_inputs + 1) + " inputs, got " + std::to_string(node.input_ids.size()));
    }

    const auto& hidden_buffer = get_input(node, 0, nodes, node_index_map);
    const auto& routing_buffer = get_input(node, 1, nodes, node_index_map);
    const auto& topk_idx_buffer = get_input(node, 2, nodes, node_index_map);

    if (hidden_buffer.precision != Precision::FP16 || node.output_buffer.precision != Precision::FP16) {
        throw std::runtime_error("moe_layer expects FP16 hidden/output");
    }
    if (topk_idx_buffer.precision != Precision::FP32) {
        throw std::runtime_error("moe_layer expects FP32 topk indices");
    }

    const __fp16* expert_scales_fp16 = nullptr;
    if (has_per_expert_scale) {
        const auto& scale_buffer = get_input(node, base_inputs, nodes, node_index_map);
        if (scale_buffer.precision != Precision::FP16) {
            throw std::runtime_error("moe_layer expects FP16 per_expert_scale");
        }
        expert_scales_fp16 = scale_buffer.data_as<__fp16>();
    }

    const size_t token_count = hidden_buffer.shape[0];
    const size_t hidden_dim = hidden_buffer.shape[1];
    const size_t total_num_experts = routing_buffer.shape[1];

    const auto& w1_0_buffer = get_input(node, 3, nodes, node_index_map);
    const size_t expert_intermediate_dim = w1_0_buffer.shape[0];

    const auto* hidden = hidden_buffer.data_as<__fp16>();
    auto* output = node.output_buffer.data_as<__fp16>();
    const auto* topk_idx = topk_idx_buffer.data_as<float>();
    const auto* routing_fp16 = routing_buffer.precision == Precision::FP16 ? routing_buffer.data_as<__fp16>() : nullptr;
    const auto* routing_fp32 = routing_buffer.precision == Precision::FP32 ? routing_buffer.data_as<float>() : nullptr;

    auto routing_prob = [&](size_t tok, size_t exp) -> float {
        const size_t offset = tok * total_num_experts + exp;
        if (routing_fp16) return static_cast<float>(routing_fp16[offset]);
        return routing_fp32[offset];
    };

    ensure_moe_buffers(token_count, hidden_dim, expert_intermediate_dim, num_experts, top_k);

    size_t* expert_offsets = moe_expert_offsets_buf.data(); 
    size_t* expert_tokens_flat = moe_expert_tokens_buf.data();