Method apply

src/fuse_attention.cpp:783–816

} // namespace

void fuse_attention::apply(module_pass_manager& mpm) const
{
    std::size_t counter = 0;

    // Fuse kv-cache attention by default
    match::find_matches(mpm, find_kv_cache_attention{.counter = &counter});
    mpm.get_module().sort();
    mpm.run_pass(dead_code_elimination{});

    // Only fuse plain attention when requested
    if(attn_enabled)
    {
        match::find_matches(mpm, find_attention{.counter = &counter});
        mpm.get_module().sort();
        mpm.run_pass(dead_code_elimination{});
    }

    std::size_t num_splits = 0;
    if(flash_decoding_num_splits.has_value())
    {
        // Use the value from the constructor (for testing)
        num_splits = *flash_decoding_num_splits;
    }
    else
    {
        // Default behavior: read from the env var (for non-test use)
        num_splits = get_num_splits();
    }
    if(num_splits > 1)
    {
        match::find_matches(mpm, find_flash_decoding{.groups = num_splits});
        mpm.run_pass(dead_code_elimination{});
    }
}

} // namespace MIGRAPHX_INLINE_NS
} // namespace migraphx