Method apply

src/fuse_attention.cpp:428–609

    }

    void apply(module_pass_manager& mpm, const match::matcher_result& r) const
    {
        auto& mm            = mpm.get_module();
        auto attn_group_ins = r.instructions["group"];
        auto* submod        = attn_group_ins->module_inputs().front();

        // TODO: for this pass of flash decoding, if LSE attn, do not do flash decoding
        auto return_ins = std::prev(submod->end());
        assert(return_ins->name() == "@return" and
               "Last instruction must be a @return instruction");
        if(return_ins->inputs().size() > 1)
            return;

        // get gemm1 and gemm2
        auto [gemm1, gemm2] = get_gemms(submod);

        // TODO: for this first pass of flash decoding, assuming no input fusion / not supporting
        auto q_param = gemm1->inputs()[0];
        auto k_param = gemm1->inputs()[1];
        auto v_param = gemm2->inputs()[1];
        assert(q_param->name() == "@param" and "Q should be a parameter");
        assert(k_param->name() == "@param" and "K should be a parameter");
        assert(v_param->name() == "@param" and "V should be a parameter");

        // check if N dimension is evenly divisible by num_splits
        if(k_param->get_shape().lens().back() % groups != 0)
            return;

        // get Q, V, K shapes from gemms
        auto qkv_shapes = get_qkv_shapes(q_param, k_param, v_param);

        // check shapes are ok and get flash decoding transformed shapes (Q', V', K')
        auto transform_info = get_transformed_shapes(qkv_shapes);

        // create mapping from submodule params to main module inputs
        auto group_inputs      = attn_group_ins->inputs();
        auto map_param_to_main = map_submod_params_to_inputs(submod, group_inputs);

        // get actual Q, K, V instructions from main module
        auto q = map_param_to_main.at(q_param);
        auto k = map_param_to_main.at(k_param);
        auto v = map_param_to_main.at(v_param);

        // insert reshape operations before group, for Q, K, V
        auto q_ndim    = q->get_shape().lens().size();
        int64_t g_axis = q_ndim - 2;

        // Q: [B, M, k] -> [B, G, M, k] via unsqueeze + broadcast
        auto q_unsqueeze =
            mm.insert_instruction(attn_group_ins, make_op("unsqueeze", {{"axes", {g_axis}}}), q);
        auto q_reshaped =
            mm.insert_instruction(attn_group_ins,
                                  make_op("multibroadcast", {{"out_lens", transform_info.q_shape}}),
                                  q_unsqueeze);

        // K: [B, k, N] -> [B, G, k, N/G] via reshape + transpose
        auto k_reshaped_intermediate = mm.insert_instruction(
            attn_group_ins, make_op("reshape", {{"dims", transform_info.k_intermediate}}), k);