Buckets:

gijl
/

lv-storage

Files

xet

gijl/lv-storage / llama.cpp /tests /test-fusion.cpp

gijl

15 days ago

download

raw

24 kB

// test-fusion: verify the backend fusion logic against a per-device baseline.
//
// for every dummy model generated by test-llama-archs, the tool runs the model on a single
// device with fusion enabled and disabled, and reports:
//   - the per-fusion-type counters for each mode (prefill / decode, merged into "any" when the
//     per-graph counts match)
//   - the NMSE between the fused and unfused logits
//   - the NMSE between the device and a CPU reference
//
// the per-fusion-type counters are compared against a per-device baseline file (CSV) so a
// fusion pattern that silently stops matching (or fires when it should not) is caught as a
// regression.
//
// usage:
//   test-fusion --models DIR --device MTL0 --record baseline.csv   # generate a baseline
//   test-fusion --models DIR --device MTL0 --check  baseline.csv   # validate against it
//   test-fusion --model FILE --device MTL0 --check  baseline.csv   # validate a single model

#include "common.h"
#include "log.h"
#include "llama-cpp.h"

#include "ggml.h"
#include "gguf.h"

#include <algorithm>
#include <array>
#include <cstring>
#include <filesystem>
#include <fstream>
#include <iomanip>
#include <iostream>
#include <map>
#include <random>
#include <set>
#include <string>
#include <vector>

// generic fusion debugging API, resolved through the ad-hoc get_proc_address mechanism
// (not part of the official ggml backend interface yet). a backend that adopts fusion debugging
// exports these exact names.
typedef void * ggml_backend_fusion_t;

typedef ggml_backend_fusion_t ( * fusion_get_t)       (ggml_backend_dev_t);
typedef void ( * fusion_stats_init_t)  (ggml_backend_fusion_t);
typedef void ( * fusion_stats_reset_t) (ggml_backend_fusion_t);
typedef int  ( * fusion_stats_get_t)   (ggml_backend_fusion_t, const char **, uint64_t *, int);
typedef void ( * fusion_set_enabled_t) (ggml_backend_fusion_t, bool);

static bool silent_model_load_progress(float, void *) {
    return true;
}

struct gguf_context_ptr {
    gguf_context * ctx;
    gguf_context_ptr(gguf_context * c) : ctx(c) {}
    ~gguf_context_ptr() { if (ctx) { gguf_free(ctx); } }
    gguf_context * get() const { return ctx; }
    gguf_context_ptr(const gguf_context_ptr &) = delete;
    gguf_context_ptr & operator=(const gguf_context_ptr &) = delete;
};

// NMSE between two vectors (same as tests/test-llama-archs.cpp)
static double nmse(const std::vector<float> & a, const std::vector<float> & b) {
    GGML_ASSERT(a.size() == b.size());
    double mse_a_b = 0.0;
    double mse_a_0 = 0.0;

    for (size_t i = 0; i < a.size(); i++) {
        const float a_i = a[i];
        const float b_i = b[i];

        mse_a_b += (a_i - b_i) * (a_i - b_i);
        mse_a_0 += a_i * a_i;
    }

    return mse_a_b / mse_a_0;
}

// deterministic token sequence
static std::vector<llama_token> get_tokens(const uint32_t n_tokens, const uint32_t n_vocab, const size_t seed) {
    std::mt19937 gen(seed);
    std::uniform_int_distribution<> dis(0, n_vocab - 1);
    std::vector<llama_token> ret;
    ret.reserve(n_tokens);
    for (uint32_t i = 0; i < n_tokens; i++) {
        ret.push_back(dis(gen));
    }
    return ret;
}

// trim leading/trailing whitespace (used when parsing padded CSV columns)
static std::string trim(const std::string & s) {
    const size_t b = s.find_first_not_of(" \t\r\n");
    if (b == std::string::npos) {
        return "";
    }
    const size_t e = s.find_last_not_of(" \t\r\n");
    return s.substr(b, e - b + 1);
}

static std::string get_arch(const std::string & path) {
    gguf_init_params params = { /*no_alloc=*/true, /*ctx=*/nullptr };
    gguf_context_ptr ctx(gguf_init_from_file(path.c_str(), params));
    if (!ctx.get()) {
        throw std::runtime_error("failed to read gguf: " + path);
    }
    const int idx = gguf_find_key(ctx.get(), "general.architecture");
    if (idx < 0) {
        return "unknown";
    }
    const char * val = gguf_get_val_str(ctx.get(), idx);
    return val ? val : "unknown";
}

static llama_model_ptr load_model(const std::string & path, ggml_backend_dev_t dev) {
    llama_model_params model_params = llama_model_default_params();
    model_params.progress_callback = silent_model_load_progress;
    std::vector<ggml_backend_dev_t> devs = { dev, nullptr };
    model_params.devices = devs.data();
    model_params.split_mode = LLAMA_SPLIT_MODE_LAYER;

    llama_model_ptr model(llama_model_load_from_file(path.c_str(), model_params));
    if (!model) {
        throw std::runtime_error("failed to load model: " + path);
    }
    return model;
}

// a fresh context (fresh state) from an already-loaded model
static llama_context_ptr create_ctx(llama_model * model, int n_ubatch) {
    llama_context_params ctx_params = llama_context_default_params();
    ctx_params.n_ctx = 0;
    ctx_params.n_threads = 4;
    ctx_params.n_threads_batch = 4;
    ctx_params.n_ubatch = n_ubatch;
    ctx_params.n_batch = n_ubatch;

    llama_context_ptr lctx(llama_init_from_model(model, ctx_params));
    if (!lctx) {
        throw std::runtime_error("failed to init context");
    }
    return lctx;
}

// decode all tokens in one batch; returns the logits of every token
static std::vector<float> decode_prefill(llama_model * model, llama_context * lctx, const std::vector<llama_token> & tokens) {
    const uint32_t n_vocab = llama_vocab_n_tokens(llama_model_get_vocab(model));
    common_batch batch(lctx);
    for (size_t i = 0; i < tokens.size(); i++) {
        batch.add(tokens[i], i, 0, true);
    }
    if (llama_process(lctx, LLAMA_PROCESS_TYPE_DECODE, batch.get())) {
        throw std::runtime_error("prefill decode failed");
    }

    std::vector<float> ret;
    ret.reserve(tokens.size() * n_vocab);
    for (size_t i = 0; i < tokens.size(); i++) {
        const float * logits_ith = llama_get_logits_ith(lctx, i);
        for (uint32_t j = 0; j < n_vocab; j++) {
            ret.push_back(logits_ith[j]);
        }
    }
    return ret;
}

// decode one token at a time; returns the logits of the last token of each step
static std::vector<float> decode_gen(llama_model * model, llama_context * lctx, const std::vector<llama_token> & tokens) {
    const uint32_t n_vocab = llama_vocab_n_tokens(llama_model_get_vocab(model));
    common_batch batch(lctx);
    std::vector<float> ret;
    for (size_t i = 0; i < tokens.size(); i++) {
        batch.clear();
        batch.add(tokens[i], i, 0, true);
        if (llama_process(lctx, LLAMA_PROCESS_TYPE_DECODE, batch.get())) {
            throw std::runtime_error("decode failed");
        }
        const float * logits = llama_get_logits_ith(lctx, 0);
        for (uint32_t j = 0; j < n_vocab; j++) {
            ret.push_back(logits[j]);
        }
    }
    return ret;
}

static void read_counts(fusion_stats_get_t api_stats_get, ggml_backend_fusion_t finfo,
                        std::vector<const char *> & labels, std::vector<uint64_t> & counts) {
    const int n = api_stats_get(finfo, nullptr, nullptr, 0);
    labels.assign(n, nullptr);
    counts.assign(n, 0);
    api_stats_get(finfo, labels.data(), counts.data(), n);
}

struct baseline_row {
    std::string arch;
    bool        moe;
    std::string mode;
    std::string label;
    uint64_t    count;
};

// one row of the per-label report
struct fusion_row {
    std::string  arch;
    bool         moe;
    std::string  mode;
    std::string  label;
    uint64_t     count_fused;
    uint64_t     count_unfused;
    uint64_t     expected;
    double       nmse_fus;
    double       nmse_dev;
    bool         ok_count; // counts match the baseline
    bool         ok_nmse;  // nmse within epsilon
};

static void usage(const char * argv0) {
    printf("%s: verify fusion counts on a device against a per-device baseline\n\n", argv0);
    printf("usage: %s [options]\n\n", argv0);
    printf("options:\n");
    printf("  --models DIR   run over all .gguf models in a directory\n");
    printf("  --model FILE   run over a single model file (mutually exclusive with --models)\n");
    printf("  --device NAME  device to run on (e.g. MTL0, CPU)\n");
    printf("  --record CSV   write the golden baseline\n");
    printf("  --check  CSV   validate the counters against a baseline (default)\n");
    printf("  -h, --help     show this message and exit\n");
}

int main(int argc, char ** argv) {
    std::string models_dir;
    std::string model_file;
    std::string device_name;
    std::string record_path;
    std::string check_path;

    for (int i = 1; i < argc; i++) {
        const std::string arg = argv[i];
        const auto next = [&](const char * name) -> std::string {
            if (i + 1 >= argc) {
                LOG_ERR("%s: %s requires an argument\n", __func__, name);
                exit(1);
            }
            return argv[++i];
        };
        if (arg == "-h" || arg == "--help") {
            usage(argv[0]);
            exit(0);
        }
        if (arg == "--models")     { models_dir  = next("--models"); }
        else if (arg == "--model") { model_file  = next("--model"); }
        else if (arg == "--device"){ device_name = next("--device"); }
        else if (arg == "--record"){ record_path = next("--record"); }
        else if (arg == "--check") { check_path  = next("--check"); }
        else {
            LOG_ERR("%s: unknown argument: %s\n", __func__, arg.c_str());
            return 1;
        }
    }

    if (device_name.empty()) {
        LOG_ERR("%s: --device NAME is required\n", __func__);
        return 1;
    }
    if (models_dir.empty() && model_file.empty()) {
        LOG_ERR("%s: --models DIR or --model FILE is required\n", __func__);
        return 1;
    }
    if (!models_dir.empty() && !model_file.empty()) {
        LOG_ERR("%s: --models DIR and --model FILE are mutually exclusive\n", __func__);
        return 1;
    }
    if (!record_path.empty() && !check_path.empty()) {
        LOG_ERR("%s: --record and --check are mutually exclusive\n", __func__);
        return 1;
    }

    std::vector<std::string> models;
    if (!model_file.empty()) {
        if (!std::filesystem::is_regular_file(model_file)) {
            LOG_ERR("%s: model file '%s' does not exist\n", __func__, model_file.c_str());
            return 1;
        }
        models.push_back(model_file);
    } else {
        if (!std::filesystem::exists(models_dir) || !std::filesystem::is_directory(models_dir)) {
            LOG_ERR("%s: models directory '%s' does not exist\n", __func__, models_dir.c_str());
            return 1;
        }
        for (const auto & entry : std::filesystem::directory_iterator(models_dir)) {
            if (entry.is_regular_file() && entry.path().extension() == ".gguf") {
                models.push_back(entry.path().string());
            }
        }
        std::sort(models.begin(), models.end());

        if (models.empty()) {
            LOG_ERR("%s: no .gguf models found in '%s'\n", __func__, models_dir.c_str());
            return 1;
        }
    }

    common_init();
    llama_backend_init();

    ggml_backend_dev_t dev = ggml_backend_dev_by_name(device_name.c_str());
    if (!dev) {
        LOG_WRN("%s: device '%s' not found - skipping (baseline is device-specific)\n",
                __func__, device_name.c_str());
        return 0;
    }

    // resolve the generic fusion debugging functions through the ad-hoc get_proc_address
    // mechanism; a backend that does not adopt fusion debugging exports none of them
    auto * reg = ggml_backend_dev_backend_reg(dev);

    // output naming uses the backend base name (e.g. "MTL") rather than the specific device
    // name (e.g. "MTL0") the test was invoked with
    const std::string base_name = ggml_backend_reg_name(reg);

    auto api_get         = (fusion_get_t)         ggml_backend_reg_get_proc_address(reg, "ggml_backend_fusion_get");
    auto api_stats_init  = (fusion_stats_init_t)  ggml_backend_reg_get_proc_address(reg, "ggml_backend_fusion_stats_init");
    auto api_stats_reset = (fusion_stats_reset_t) ggml_backend_reg_get_proc_address(reg, "ggml_backend_fusion_stats_reset");
    auto api_stats_get   = (fusion_stats_get_t)   ggml_backend_reg_get_proc_address(reg, "ggml_backend_fusion_stats_get");
    auto api_set_enabled = (fusion_set_enabled_t) ggml_backend_reg_get_proc_address(reg, "ggml_backend_fusion_set_enabled");

    if (!api_get || !api_stats_init || !api_set_enabled || !api_stats_reset || !api_stats_get) {
        LOG_ERR("%s: device '%s' does not export the generic fusion debugging API "
                "(ggml_backend_fusion_*) - cannot run the fusion regression test\n",
                __func__, device_name.c_str());
        return 1;
    }

    ggml_backend_fusion_t finfo = api_get(dev);

    // enable fusions stats
    api_stats_init(finfo);

    const bool has_counts = true;

    // load the baseline (if any): key arch|moe|mode|label -> expected row
    std::map<std::string, baseline_row> baseline;
    if (!check_path.empty()) {
        std::ifstream in(check_path);
        if (!in) {
            LOG_ERR("%s: cannot open baseline '%s'\n", __func__, check_path.c_str());
            return 1;
        }
        std::string line;
        while (std::getline(in, line)) {
            if (line.empty() || line[0] == '#') {
                continue;
            }
            std::vector<std::string> cols;
            size_t pos = 0;
            while ((pos = line.find(',')) != std::string::npos) {
                cols.push_back(trim(line.substr(0, pos)));
                line.erase(0, pos + 1);
            }
            cols.push_back(trim(line));
            if (cols.size() != 5) {
                continue;
            }
            const std::string key = cols[0] + "|" + cols[1] + "|" + cols[2] + "|" + cols[3];
            baseline[key] = { cols[0], cols[1] == "1", cols[2], cols[3], std::stoull(cols[4]) };
        }
    }

    std::vector<fusion_row> rows;

    std::set<std::string> run_arches;
    std::set<std::string> produced;

    LOG_INF("%s: running fusion test over %zu models on '%s'\n", __func__, models.size(), base_name.c_str());

    const size_t seed = 1;

    for (const auto & model_path : models) {
        const std::string arch = get_arch(model_path);
        const bool moe = arch.find("moe") != std::string::npos;

        llama_model_ptr model;
        llama_model_ptr model_cpu;
        uint32_t n_vocab = 0;
        try {
            model = load_model(model_path, dev);
            model_cpu = load_model(model_path, ggml_backend_dev_by_name("CPU"));
            n_vocab = llama_vocab_n_tokens(llama_model_get_vocab(model.get()));
        } catch (const std::exception & e) {
            LOG_ERR("%s: %s: %s\n", __func__, model_path.c_str(), e.what());
            continue;
        }

        run_arches.insert(arch);

        struct mode_cfg {
            std::string name;
            std::vector<float> (*decode)(llama_model *, llama_context *, const std::vector<llama_token> &);
            int n_tokens;
            int n_graphs;   // graph runs per mode (prefill=1, decode=16)
        };
        const mode_cfg modes[] = {
            { "prefill", decode_prefill, 32, 1  },
            { "decode",  decode_gen,     16, 16 },
        };

        // per-label, per-mode data for this model; prefill and decode are merged into a single
        // "any" row when their per-graph counts match
        struct mode_data {
            bool     present;
            uint64_t count_fused;   // per graph
            uint64_t count_unfused; // per graph
            double   nmse_fus;
            double   nmse_dev;
            bool     ok_nmse;
        };
        std::map<std::string, std::array<mode_data, 2>> mdata;

        for (int mi = 0; mi < 2; mi++) {
            const mode_cfg & mode = modes[mi];
            const auto tokens = get_tokens(mode.n_tokens, n_vocab, seed);

            // CPU reference for this mode (fresh context, fresh state)
            std::vector<float> logits_cpu;
            try {
                llama_context_ptr ctx = create_ctx(model_cpu.get(), 32);
                logits_cpu = mode.decode(model_cpu.get(), ctx.get(), tokens);
            } catch (const std::exception & e) {
                LOG_WRN("%s: %s: cpu reference: %s\n", __func__, model_path.c_str(), e.what());
            }

            // fused run on a fresh context (fresh state)
            std::vector<float> logits_fused;
            std::vector<const char *> labels;
            std::vector<uint64_t> counts_fused;
            {
                llama_context_ptr ctx = create_ctx(model.get(), 32);
                if (has_counts) {
                    api_set_enabled(finfo, true);
                    api_stats_reset(finfo);
                }
                logits_fused = mode.decode(model.get(), ctx.get(), tokens);
                if (has_counts) {
                    read_counts(api_stats_get, finfo, labels, counts_fused);
                }
            }

            // unfused run on another fresh context (fresh state)
            std::vector<float> logits_unfused;
            std::vector<uint64_t> counts_unfused;
            {
                llama_context_ptr ctx = create_ctx(model.get(), 32);
                if (has_counts) {
                    api_set_enabled(finfo, false);
                    api_stats_reset(finfo);
                }
                logits_unfused = mode.decode(model.get(), ctx.get(), tokens);
                if (has_counts) {
                    read_counts(api_stats_get, finfo, labels, counts_unfused);
                }
            }

            const double nmse_fus = nmse(logits_fused, logits_unfused);
            const double nmse_dev = logits_cpu.empty() ? 0.0 : nmse(logits_fused, logits_cpu);

            if (has_counts) {
                for (int i = 0; i < (int) labels.size(); i++) {
                    const uint64_t fused   = counts_fused[i]   / mode.n_graphs;
                    const uint64_t unfused = counts_unfused[i] / mode.n_graphs;
                    if (fused == 0 && unfused == 0) {
                        continue;
                    }
                    auto & d = mdata[labels[i]][mi];
                    d.present       = true;
                    d.count_fused   = fused;
                    d.count_unfused = unfused;
                    d.nmse_fus      = nmse_fus;
                    d.nmse_dev      = nmse_dev;
                    d.ok_nmse       = nmse_fus <= 1e-4;
                }
            } else {
                rows.push_back({ arch, moe, mode.name, "?", 0, 0, 0, nmse_fus, nmse_dev, true, nmse_fus <= 1e-4 });
            }
        }

        // build the per-label rows, merging prefill and decode into "any" when the per-graph
        // counts match (they always do for the deterministic fusion table)
        if (has_counts) {
            for (auto & kv : mdata) {
                const std::string & label = kv.first;
                const auto & d = kv.second;
                const bool both = d[0].present && d[1].present;
                const bool match = both && d[0].count_fused == d[1].count_fused;

                if (match) {
                    // one "any" row; use the worst NMSE across the two modes
                    const std::string any_key = arch + "|" + (moe ? "1" : "0") + "|any|" + label;
                    produced.insert(any_key);
                    const uint64_t expected = baseline.count(any_key) ? baseline.at(any_key).count : 0;
                    const bool ok_count = check_path.empty() || d[0].count_fused == expected;
                    const bool ok_nmse   = d[0].ok_nmse && d[1].ok_nmse;
                    const double nmse_fus = std::max(d[0].nmse_fus, d[1].nmse_fus);
                    const double nmse_dev = std::max(d[0].nmse_dev, d[1].nmse_dev);
                    rows.push_back({ arch, moe, "any", label, d[0].count_fused, d[0].count_unfused,
                                     expected, nmse_fus, nmse_dev, ok_count, ok_nmse });
                } else {
                    // counts differ - keep a separate row per mode
                    for (int mi = 0; mi < 2; mi++) {
                        if (!d[mi].present) {
                            continue;
                        }
                        const mode_data & a = d[mi];
                        const std::string mode_key = arch + "|" + (moe ? "1" : "0") + "|" + modes[mi].name + "|" + label;
                        produced.insert(mode_key);
                        const uint64_t expected = baseline.count(mode_key) ? baseline.at(mode_key).count : 0;
                        const bool ok_count = check_path.empty() || a.count_fused == expected;
                        rows.push_back({ arch, moe, modes[mi].name, label, a.count_fused, a.count_unfused,
                                         expected, a.nmse_fus, a.nmse_dev, ok_count, a.ok_nmse });
                    }
                }
            }
        }

        LOG_INF("%s: %-20s (%s) done\n", __func__, arch.c_str(), model_path.c_str());
    }

    if (!check_path.empty()) {
        for (const auto & kv : baseline) {
            if (produced.count(kv.first) || !run_arches.count(kv.second.arch)) {
                continue;
            }
            const baseline_row & b = kv.second;
            LOG_WRN("%s: %s: fusion %s x%llu is no longer matched\n", __func__, b.arch.c_str(),
                    b.label.c_str(), (unsigned long long) b.count);
            rows.push_back({ b.arch, b.moe, b.mode, b.label, 0, 0, b.count, 0.0, 0.0, false, true });
        }
    }

    // print the report
    {
        std::ofstream out(record_path);
        std::ostream & os = record_path.empty() ? std::cout : out;
        if (!record_path.empty()) {
            os << "# test-fusion baseline for device " << base_name << "\n";
            os << "# " << std::left
               << std::setw(18) << "arch"  << ','
               << std::setw(4)  << "moe"   << ','
               << std::setw(8)  << "mode"  << ','
               << std::setw(28) << "label" << ','
               << std::right << std::setw(7) << "count" << '\n';
        }

        LOG_INF("%-20s %-4s %-8s %-22s %7s %7s %7s %10s %10s %s\n",
                "arch", "moe", "mode", "label", "fused", "unfused", "expected", "nmse_fus", "nmse_dev", "status");
        int n_ok = 0;
        int n_bad = 0;
        for (const auto & r : rows) {
            const bool ok = r.ok_count && r.ok_nmse;
            const char * status = ok ? "ok" : "FAIL";
            if (ok) { n_ok++; } else { n_bad++; }
            LOG_INF("%-20s %-4s %-8s %-22s %7llu %7llu %7llu %10.2e %10.2e %s\n",
                    r.arch.c_str(), r.moe ? "moe" : "dense", r.mode.c_str(), r.label.c_str(),
                    (unsigned long long) r.count_fused, (unsigned long long) r.count_unfused,
                    (unsigned long long) r.expected, r.nmse_fus, r.nmse_dev, status);
            if (!record_path.empty()) {
                os << std::left
                   << std::setw(20) << r.arch << ','
                   << std::setw(4)  << (r.moe ? "1" : "0") << ','
                   << std::setw(8)  << r.mode << ','
                   << std::setw(28) << r.label << ','
                   << std::right << std::setw(7) << r.count_fused << '\n';
            }
        }
        LOG_INF("summary: %d ok, %d failed\n", n_ok, n_bad);
        if (!record_path.empty()) {
            LOG_INF("%s: baseline written to '%s'\n", __func__, record_path.c_str());
        }

        if (n_bad && !models_dir.empty() && !check_path.empty()) {
            LOG_WRN("%s: if the fusion counts are expected to change, run with --record to update the baseline:\n"
                    "\n"
                    "./bin/test-llama-archs -o %s\n"
                    "%s --device %s --models %s --record %s\n",
                    __func__, models_dir.c_str(), argv[0], device_name.c_str(), models_dir.c_str(), check_path.c_str());
        }

        return n_bad;
    }
}

Xet Storage Details

Size:: 24 kB
Xet hash:: 0d77acd27f952d0e5f30214808c073c044e578da4c1c5e223fb46f623f7c4db2

Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.

	// test-fusion: verify the backend fusion logic against a per-device baseline.
	//
	// for every dummy model generated by test-llama-archs, the tool runs the model on a single
	// device with fusion enabled and disabled, and reports:
	// - the per-fusion-type counters for each mode (prefill / decode, merged into "any" when the
	// per-graph counts match)
	// - the NMSE between the fused and unfused logits
	// - the NMSE between the device and a CPU reference
	//
	// the per-fusion-type counters are compared against a per-device baseline file (CSV) so a
	// fusion pattern that silently stops matching (or fires when it should not) is caught as a
	// regression.
	//
	// usage:
	// test-fusion --models DIR --device MTL0 --record baseline.csv # generate a baseline
	// test-fusion --models DIR --device MTL0 --check baseline.csv # validate against it
	// test-fusion --model FILE --device MTL0 --check baseline.csv # validate a single model

	#include "common.h"
	#include "log.h"
	#include "llama-cpp.h"

	#include "ggml.h"
	#include "gguf.h"

	#include <algorithm>
	#include <array>
	#include <cstring>
	#include <filesystem>
	#include <fstream>
	#include <iomanip>
	#include <iostream>
	#include <map>
	#include <random>
	#include <set>
	#include <string>
	#include <vector>

	// generic fusion debugging API, resolved through the ad-hoc get_proc_address mechanism
	// (not part of the official ggml backend interface yet). a backend that adopts fusion debugging
	// exports these exact names.
	typedef void * ggml_backend_fusion_t;

	typedef ggml_backend_fusion_t ( * fusion_get_t) (ggml_backend_dev_t);
	typedef void ( * fusion_stats_init_t) (ggml_backend_fusion_t);
	typedef void ( * fusion_stats_reset_t) (ggml_backend_fusion_t);
	typedef int ( * fusion_stats_get_t) (ggml_backend_fusion_t, const char *, uint64_t , int);
	typedef void ( * fusion_set_enabled_t) (ggml_backend_fusion_t, bool);

	static bool silent_model_load_progress(float, void *) {
	return true;
	}

	struct gguf_context_ptr {
	gguf_context * ctx;
	gguf_context_ptr(gguf_context * c) : ctx(c) {}
	~gguf_context_ptr() { if (ctx) { gguf_free(ctx); } }
	gguf_context * get() const { return ctx; }
	gguf_context_ptr(const gguf_context_ptr &) = delete;
	gguf_context_ptr & operator=(const gguf_context_ptr &) = delete;
	};

	// NMSE between two vectors (same as tests/test-llama-archs.cpp)
	static double nmse(const std::vector<float> & a, const std::vector<float> & b) {
	GGML_ASSERT(a.size() == b.size());
	double mse_a_b = 0.0;
	double mse_a_0 = 0.0;

	for (size_t i = 0; i < a.size(); i++) {
	const float a_i = a[i];
	const float b_i = b[i];

	mse_a_b += (a_i - b_i) * (a_i - b_i);
	mse_a_0 += a_i * a_i;
	}

	return mse_a_b / mse_a_0;
	}

	// deterministic token sequence
	static std::vector<llama_token> get_tokens(const uint32_t n_tokens, const uint32_t n_vocab, const size_t seed) {
	std::mt19937 gen(seed);
	std::uniform_int_distribution<> dis(0, n_vocab - 1);
	std::vector<llama_token> ret;
	ret.reserve(n_tokens);
	for (uint32_t i = 0; i < n_tokens; i++) {
	ret.push_back(dis(gen));
	}
	return ret;
	}

	// trim leading/trailing whitespace (used when parsing padded CSV columns)
	static std::string trim(const std::string & s) {
	const size_t b = s.find_first_not_of(" \t\r\n");
	if (b == std::string::npos) {
	return "";
	}
	const size_t e = s.find_last_not_of(" \t\r\n");
	return s.substr(b, e - b + 1);
	}

	static std::string get_arch(const std::string & path) {
	gguf_init_params params = { /no_alloc=/true, /ctx=/nullptr };
	gguf_context_ptr ctx(gguf_init_from_file(path.c_str(), params));
	if (!ctx.get()) {
	throw std::runtime_error("failed to read gguf: " + path);
	}
	const int idx = gguf_find_key(ctx.get(), "general.architecture");
	if (idx < 0) {
	return "unknown";
	}
	const char * val = gguf_get_val_str(ctx.get(), idx);
	return val ? val : "unknown";
	}

	static llama_model_ptr load_model(const std::string & path, ggml_backend_dev_t dev) {
	llama_model_params model_params = llama_model_default_params();
	model_params.progress_callback = silent_model_load_progress;
	std::vector<ggml_backend_dev_t> devs = { dev, nullptr };
	model_params.devices = devs.data();
	model_params.split_mode = LLAMA_SPLIT_MODE_LAYER;

	llama_model_ptr model(llama_model_load_from_file(path.c_str(), model_params));
	if (!model) {
	throw std::runtime_error("failed to load model: " + path);
	}
	return model;
	}

	// a fresh context (fresh state) from an already-loaded model
	static llama_context_ptr create_ctx(llama_model * model, int n_ubatch) {
	llama_context_params ctx_params = llama_context_default_params();
	ctx_params.n_ctx = 0;
	ctx_params.n_threads = 4;
	ctx_params.n_threads_batch = 4;
	ctx_params.n_ubatch = n_ubatch;
	ctx_params.n_batch = n_ubatch;

	llama_context_ptr lctx(llama_init_from_model(model, ctx_params));
	if (!lctx) {
	throw std::runtime_error("failed to init context");
	}
	return lctx;
	}

	// decode all tokens in one batch; returns the logits of every token
	static std::vector<float> decode_prefill(llama_model * model, llama_context * lctx, const std::vector<llama_token> & tokens) {
	const uint32_t n_vocab = llama_vocab_n_tokens(llama_model_get_vocab(model));
	common_batch batch(lctx);
	for (size_t i = 0; i < tokens.size(); i++) {
	batch.add(tokens[i], i, 0, true);
	}
	if (llama_process(lctx, LLAMA_PROCESS_TYPE_DECODE, batch.get())) {
	throw std::runtime_error("prefill decode failed");
	}

	std::vector<float> ret;
	ret.reserve(tokens.size() * n_vocab);
	for (size_t i = 0; i < tokens.size(); i++) {
	const float * logits_ith = llama_get_logits_ith(lctx, i);
	for (uint32_t j = 0; j < n_vocab; j++) {
	ret.push_back(logits_ith[j]);
	}
	}
	return ret;
	}

	// decode one token at a time; returns the logits of the last token of each step
	static std::vector<float> decode_gen(llama_model * model, llama_context * lctx, const std::vector<llama_token> & tokens) {
	const uint32_t n_vocab = llama_vocab_n_tokens(llama_model_get_vocab(model));
	common_batch batch(lctx);
	std::vector<float> ret;
	for (size_t i = 0; i < tokens.size(); i++) {
	batch.clear();
	batch.add(tokens[i], i, 0, true);
	if (llama_process(lctx, LLAMA_PROCESS_TYPE_DECODE, batch.get())) {
	throw std::runtime_error("decode failed");
	}
	const float * logits = llama_get_logits_ith(lctx, 0);
	for (uint32_t j = 0; j < n_vocab; j++) {
	ret.push_back(logits[j]);
	}
	}
	return ret;
	}

	static void read_counts(fusion_stats_get_t api_stats_get, ggml_backend_fusion_t finfo,
	std::vector<const char *> & labels, std::vector<uint64_t> & counts) {
	const int n = api_stats_get(finfo, nullptr, nullptr, 0);
	labels.assign(n, nullptr);
	counts.assign(n, 0);
	api_stats_get(finfo, labels.data(), counts.data(), n);
	}

	struct baseline_row {
	std::string arch;
	bool moe;
	std::string mode;
	std::string label;
	uint64_t count;
	};

	// one row of the per-label report
	struct fusion_row {
	std::string arch;
	bool moe;
	std::string mode;
	std::string label;
	uint64_t count_fused;
	uint64_t count_unfused;
	uint64_t expected;
	double nmse_fus;
	double nmse_dev;
	bool ok_count; // counts match the baseline
	bool ok_nmse; // nmse within epsilon
	};

	static void usage(const char * argv0) {
	printf("%s: verify fusion counts on a device against a per-device baseline\n\n", argv0);
	printf("usage: %s [options]\n\n", argv0);
	printf("options:\n");
	printf(" --models DIR run over all .gguf models in a directory\n");
	printf(" --model FILE run over a single model file (mutually exclusive with --models)\n");
	printf(" --device NAME device to run on (e.g. MTL0, CPU)\n");
	printf(" --record CSV write the golden baseline\n");
	printf(" --check CSV validate the counters against a baseline (default)\n");
	printf(" -h, --help show this message and exit\n");
	}

	int main(int argc, char ** argv) {
	std::string models_dir;
	std::string model_file;
	std::string device_name;
	std::string record_path;
	std::string check_path;

	for (int i = 1; i < argc; i++) {
	const std::string arg = argv[i];
	const auto next = [&](const char * name) -> std::string {
	if (i + 1 >= argc) {
	LOG_ERR("%s: %s requires an argument\n", __func__, name);
	exit(1);
	}
	return argv[++i];
	};
	if (arg == "-h" \|\| arg == "--help") {
	usage(argv[0]);
	exit(0);
	}
	if (arg == "--models") { models_dir = next("--models"); }
	else if (arg == "--model") { model_file = next("--model"); }
	else if (arg == "--device"){ device_name = next("--device"); }
	else if (arg == "--record"){ record_path = next("--record"); }
	else if (arg == "--check") { check_path = next("--check"); }
	else {
	LOG_ERR("%s: unknown argument: %s\n", __func__, arg.c_str());
	return 1;
	}
	}

	if (device_name.empty()) {
	LOG_ERR("%s: --device NAME is required\n", __func__);
	return 1;
	}
	if (models_dir.empty() && model_file.empty()) {
	LOG_ERR("%s: --models DIR or --model FILE is required\n", __func__);
	return 1;
	}
	if (!models_dir.empty() && !model_file.empty()) {
	LOG_ERR("%s: --models DIR and --model FILE are mutually exclusive\n", __func__);
	return 1;
	}
	if (!record_path.empty() && !check_path.empty()) {
	LOG_ERR("%s: --record and --check are mutually exclusive\n", __func__);
	return 1;
	}

	std::vector<std::string> models;
	if (!model_file.empty()) {
	if (!std::filesystem::is_regular_file(model_file)) {
	LOG_ERR("%s: model file '%s' does not exist\n", __func__, model_file.c_str());
	return 1;
	}
	models.push_back(model_file);
	} else {
	if (!std::filesystem::exists(models_dir) \|\| !std::filesystem::is_directory(models_dir)) {
	LOG_ERR("%s: models directory '%s' does not exist\n", __func__, models_dir.c_str());
	return 1;
	}
	for (const auto & entry : std::filesystem::directory_iterator(models_dir)) {
	if (entry.is_regular_file() && entry.path().extension() == ".gguf") {
	models.push_back(entry.path().string());
	}
	}
	std::sort(models.begin(), models.end());

	if (models.empty()) {
	LOG_ERR("%s: no .gguf models found in '%s'\n", __func__, models_dir.c_str());
	return 1;
	}
	}

	common_init();
	llama_backend_init();

	ggml_backend_dev_t dev = ggml_backend_dev_by_name(device_name.c_str());
	if (!dev) {
	LOG_WRN("%s: device '%s' not found - skipping (baseline is device-specific)\n",
	__func__, device_name.c_str());
	return 0;
	}

	// resolve the generic fusion debugging functions through the ad-hoc get_proc_address
	// mechanism; a backend that does not adopt fusion debugging exports none of them
	auto * reg = ggml_backend_dev_backend_reg(dev);

	// output naming uses the backend base name (e.g. "MTL") rather than the specific device
	// name (e.g. "MTL0") the test was invoked with
	const std::string base_name = ggml_backend_reg_name(reg);

	auto api_get = (fusion_get_t) ggml_backend_reg_get_proc_address(reg, "ggml_backend_fusion_get");
	auto api_stats_init = (fusion_stats_init_t) ggml_backend_reg_get_proc_address(reg, "ggml_backend_fusion_stats_init");
	auto api_stats_reset = (fusion_stats_reset_t) ggml_backend_reg_get_proc_address(reg, "ggml_backend_fusion_stats_reset");
	auto api_stats_get = (fusion_stats_get_t) ggml_backend_reg_get_proc_address(reg, "ggml_backend_fusion_stats_get");
	auto api_set_enabled = (fusion_set_enabled_t) ggml_backend_reg_get_proc_address(reg, "ggml_backend_fusion_set_enabled");

	if (!api_get \|\| !api_stats_init \|\| !api_set_enabled \|\| !api_stats_reset \|\| !api_stats_get) {
	LOG_ERR("%s: device '%s' does not export the generic fusion debugging API "
	"(ggml_backend_fusion_*) - cannot run the fusion regression test\n",
	__func__, device_name.c_str());
	return 1;
	}

	ggml_backend_fusion_t finfo = api_get(dev);

	// enable fusions stats
	api_stats_init(finfo);

	const bool has_counts = true;

	// load the baseline (if any): key arch\|moe\|mode\|label -> expected row
	std::map<std::string, baseline_row> baseline;
	if (!check_path.empty()) {
	std::ifstream in(check_path);
	if (!in) {
	LOG_ERR("%s: cannot open baseline '%s'\n", __func__, check_path.c_str());
	return 1;
	}
	std::string line;
	while (std::getline(in, line)) {
	if (line.empty() \|\| line[0] == '#') {
	continue;
	}
	std::vector<std::string> cols;
	size_t pos = 0;
	while ((pos = line.find(',')) != std::string::npos) {
	cols.push_back(trim(line.substr(0, pos)));
	line.erase(0, pos + 1);
	}
	cols.push_back(trim(line));
	if (cols.size() != 5) {
	continue;
	}
	const std::string key = cols[0] + "\|" + cols[1] + "\|" + cols[2] + "\|" + cols[3];
	baseline[key] = { cols[0], cols[1] == "1", cols[2], cols[3], std::stoull(cols[4]) };
	}
	}

	std::vector<fusion_row> rows;

	std::set<std::string> run_arches;
	std::set<std::string> produced;

	LOG_INF("%s: running fusion test over %zu models on '%s'\n", __func__, models.size(), base_name.c_str());

	const size_t seed = 1;

	for (const auto & model_path : models) {
	const std::string arch = get_arch(model_path);
	const bool moe = arch.find("moe") != std::string::npos;

	llama_model_ptr model;
	llama_model_ptr model_cpu;
	uint32_t n_vocab = 0;
	try {
	model = load_model(model_path, dev);
	model_cpu = load_model(model_path, ggml_backend_dev_by_name("CPU"));
	n_vocab = llama_vocab_n_tokens(llama_model_get_vocab(model.get()));
	} catch (const std::exception & e) {
	LOG_ERR("%s: %s: %s\n", __func__, model_path.c_str(), e.what());
	continue;
	}

	run_arches.insert(arch);

	struct mode_cfg {
	std::string name;
	std::vector<float> (decode)(llama_model , llama_context *, const std::vector<llama_token> &);
	int n_tokens;
	int n_graphs; // graph runs per mode (prefill=1, decode=16)
	};
	const mode_cfg modes[] = {
	{ "prefill", decode_prefill, 32, 1 },
	{ "decode", decode_gen, 16, 16 },
	};

	// per-label, per-mode data for this model; prefill and decode are merged into a single
	// "any" row when their per-graph counts match
	struct mode_data {
	bool present;
	uint64_t count_fused; // per graph
	uint64_t count_unfused; // per graph
	double nmse_fus;
	double nmse_dev;
	bool ok_nmse;
	};
	std::map<std::string, std::array<mode_data, 2>> mdata;

	for (int mi = 0; mi < 2; mi++) {
	const mode_cfg & mode = modes[mi];
	const auto tokens = get_tokens(mode.n_tokens, n_vocab, seed);

	// CPU reference for this mode (fresh context, fresh state)
	std::vector<float> logits_cpu;
	try {
	llama_context_ptr ctx = create_ctx(model_cpu.get(), 32);
	logits_cpu = mode.decode(model_cpu.get(), ctx.get(), tokens);
	} catch (const std::exception & e) {
	LOG_WRN("%s: %s: cpu reference: %s\n", __func__, model_path.c_str(), e.what());
	}

	// fused run on a fresh context (fresh state)
	std::vector<float> logits_fused;
	std::vector<const char *> labels;
	std::vector<uint64_t> counts_fused;
	{
	llama_context_ptr ctx = create_ctx(model.get(), 32);
	if (has_counts) {
	api_set_enabled(finfo, true);
	api_stats_reset(finfo);
	}
	logits_fused = mode.decode(model.get(), ctx.get(), tokens);
	if (has_counts) {
	read_counts(api_stats_get, finfo, labels, counts_fused);
	}
	}

	// unfused run on another fresh context (fresh state)
	std::vector<float> logits_unfused;
	std::vector<uint64_t> counts_unfused;
	{
	llama_context_ptr ctx = create_ctx(model.get(), 32);
	if (has_counts) {
	api_set_enabled(finfo, false);
	api_stats_reset(finfo);
	}
	logits_unfused = mode.decode(model.get(), ctx.get(), tokens);
	if (has_counts) {
	read_counts(api_stats_get, finfo, labels, counts_unfused);
	}
	}

	const double nmse_fus = nmse(logits_fused, logits_unfused);
	const double nmse_dev = logits_cpu.empty() ? 0.0 : nmse(logits_fused, logits_cpu);

	if (has_counts) {
	for (int i = 0; i < (int) labels.size(); i++) {
	const uint64_t fused = counts_fused[i] / mode.n_graphs;
	const uint64_t unfused = counts_unfused[i] / mode.n_graphs;
	if (fused == 0 && unfused == 0) {
	continue;
	}
	auto & d = mdata[labels[i]][mi];
	d.present = true;
	d.count_fused = fused;
	d.count_unfused = unfused;
	d.nmse_fus = nmse_fus;
	d.nmse_dev = nmse_dev;
	d.ok_nmse = nmse_fus <= 1e-4;
	}
	} else {
	rows.push_back({ arch, moe, mode.name, "?", 0, 0, 0, nmse_fus, nmse_dev, true, nmse_fus <= 1e-4 });
	}
	}

	// build the per-label rows, merging prefill and decode into "any" when the per-graph
	// counts match (they always do for the deterministic fusion table)
	if (has_counts) {
	for (auto & kv : mdata) {
	const std::string & label = kv.first;
	const auto & d = kv.second;
	const bool both = d[0].present && d[1].present;
	const bool match = both && d[0].count_fused == d[1].count_fused;

	if (match) {
	// one "any" row; use the worst NMSE across the two modes
	const std::string any_key = arch + "\|" + (moe ? "1" : "0") + "\|any\|" + label;
	produced.insert(any_key);
	const uint64_t expected = baseline.count(any_key) ? baseline.at(any_key).count : 0;
	const bool ok_count = check_path.empty() \|\| d[0].count_fused == expected;
	const bool ok_nmse = d[0].ok_nmse && d[1].ok_nmse;
	const double nmse_fus = std::max(d[0].nmse_fus, d[1].nmse_fus);
	const double nmse_dev = std::max(d[0].nmse_dev, d[1].nmse_dev);
	rows.push_back({ arch, moe, "any", label, d[0].count_fused, d[0].count_unfused,
	expected, nmse_fus, nmse_dev, ok_count, ok_nmse });
	} else {
	// counts differ - keep a separate row per mode
	for (int mi = 0; mi < 2; mi++) {
	if (!d[mi].present) {
	continue;
	}
	const mode_data & a = d[mi];
	const std::string mode_key = arch + "\|" + (moe ? "1" : "0") + "\|" + modes[mi].name + "\|" + label;
	produced.insert(mode_key);
	const uint64_t expected = baseline.count(mode_key) ? baseline.at(mode_key).count : 0;
	const bool ok_count = check_path.empty() \|\| a.count_fused == expected;
	rows.push_back({ arch, moe, modes[mi].name, label, a.count_fused, a.count_unfused,
	expected, a.nmse_fus, a.nmse_dev, ok_count, a.ok_nmse });
	}
	}
	}
	}

	LOG_INF("%s: %-20s (%s) done\n", __func__, arch.c_str(), model_path.c_str());
	}

	if (!check_path.empty()) {
	for (const auto & kv : baseline) {
	if (produced.count(kv.first) \|\| !run_arches.count(kv.second.arch)) {
	continue;
	}
	const baseline_row & b = kv.second;
	LOG_WRN("%s: %s: fusion %s x%llu is no longer matched\n", __func__, b.arch.c_str(),
	b.label.c_str(), (unsigned long long) b.count);
	rows.push_back({ b.arch, b.moe, b.mode, b.label, 0, 0, b.count, 0.0, 0.0, false, true });
	}
	}

	// print the report
	{
	std::ofstream out(record_path);
	std::ostream & os = record_path.empty() ? std::cout : out;
	if (!record_path.empty()) {
	os << "# test-fusion baseline for device " << base_name << "\n";
	os << "# " << std::left
	<< std::setw(18) << "arch" << ','
	<< std::setw(4) << "moe" << ','
	<< std::setw(8) << "mode" << ','
	<< std::setw(28) << "label" << ','
	<< std::right << std::setw(7) << "count" << '\n';
	}

	LOG_INF("%-20s %-4s %-8s %-22s %7s %7s %7s %10s %10s %s\n",
	"arch", "moe", "mode", "label", "fused", "unfused", "expected", "nmse_fus", "nmse_dev", "status");
	int n_ok = 0;
	int n_bad = 0;
	for (const auto & r : rows) {
	const bool ok = r.ok_count && r.ok_nmse;
	const char * status = ok ? "ok" : "FAIL";
	if (ok) { n_ok++; } else { n_bad++; }
	LOG_INF("%-20s %-4s %-8s %-22s %7llu %7llu %7llu %10.2e %10.2e %s\n",
	r.arch.c_str(), r.moe ? "moe" : "dense", r.mode.c_str(), r.label.c_str(),
	(unsigned long long) r.count_fused, (unsigned long long) r.count_unfused,
	(unsigned long long) r.expected, r.nmse_fus, r.nmse_dev, status);
	if (!record_path.empty()) {
	os << std::left
	<< std::setw(20) << r.arch << ','
	<< std::setw(4) << (r.moe ? "1" : "0") << ','
	<< std::setw(8) << r.mode << ','
	<< std::setw(28) << r.label << ','
	<< std::right << std::setw(7) << r.count_fused << '\n';
	}
	}
	LOG_INF("summary: %d ok, %d failed\n", n_ok, n_bad);
	if (!record_path.empty()) {
	LOG_INF("%s: baseline written to '%s'\n", __func__, record_path.c_str());
	}

	if (n_bad && !models_dir.empty() && !check_path.empty()) {
	LOG_WRN("%s: if the fusion counts are expected to change, run with --record to update the baseline:\n"
	"\n"
	"./bin/test-llama-archs -o %s\n"
	"%s --device %s --models %s --record %s\n",
	__func__, models_dir.c_str(), argv[0], device_name.c_str(), models_dir.c_str(), check_path.c_str());
	}

	return n_bad;
	}
	}