// metal/tests/moe_mm_parity.cpp - the prompt-chunk routed experts (moe_mm.metal): plan, gate/up GEMM, down GEMM,
// combine, against a CPU reference that rounds the activations and weights to fp16 as the kernels do (the matrix
// units' operands) - and against the exact float reference for the size of that rounding; then the speed.
//
//   moe_mm_parity [--gu Q2_0] [--dn Q2_0] [--T 256] [--tile 32,64,64,4]          parity (64 experts)
//   moe_mm_parity --gather-bench [T=2048] [experts=512] [width=2560]           short plan/copy sweep
//   MM_BASELINE_KERNEL=/path/to/moe_mm.metal compares every plan buffer byte-for-byte
//   moe_mm_parity --bench [--gu ..] [--dn ..] [--T 2048] [--tile ..]            per-layer time, 512 experts
#include "coop_probe.h"
#include "mtl_rt.h"
#include "ref_quant.h"

#include <algorithm>
#include <cmath>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <set>
#include <string>
#include <vector>

namespace {
constexpr int N_EMBD = 2560, N_FF = 640, TOPK = 10;
struct MmPlanArgs { uint32_t n_ent, topk, n_expert, tm, n_blocks; };
struct MmArgs { uint64_t gate_off, up_off, down_off, gu_stride, d_stride; uint32_t gu_row, d_row, n_embd, n_ff; };
double median(std::vector<double> v) { std::sort(v.begin(), v.end()); return v[v.size() / 2]; }
int type_of(const std::string& s) {
    for (int t : {ref::Q2_0, ref::IQ3_XXS, ref::IQ2_XXS, ref::IQ2_XS, ref::IQ2_S, ref::IQ3_S, ref::IQ4_NL, ref::IQ4_XS, ref::Q4_K, ref::Q5_K, ref::Q6_K, ref::Q8_0, ref::Q5_1})
        if (s == ref::type_name(t)) return t;
    std::exit(2);
}
float h16(float v) { return ref::f16_to_f32(ref::f32_to_f16(v)); }

int gather_bench(int T, int NE, int n) {
    if (T <= 0 || NE < TOPK || NE > 1024 || NE % 32 || n <= 0 || n % 8) {
        std::fprintf(stderr, "gather-bench needs T > 0, experts in [32, 1024] divisible by 32, and width divisible by 8\n");
        return 2;
    }
    mtl::Context cx;
    const int lib = cx.compile({"moe_mm.metal"}, "#define NO_GGML_COMMON 1\n#define MM_GATHER_ONLY 1\n");
    mtl::Pipeline phist = cx.pipeline(lib, "mm_hist"), pscan = cx.pipeline(lib, "mm_scan"), pplace = cx.pipeline(lib, "mm_place"),
                  pgat = cx.pipeline(lib, "mm_gather_h"), pgat_f = cx.pipeline(lib, "mm_gather");
    const char* baseline = std::getenv("MM_BASELINE_KERNEL");
    const int bl = baseline ? cx.compile({baseline}, "#define NO_GGML_COMMON 1\n#define MM_GATHER_ONLY 1\n") : lib;
    mtl::Pipeline bh = cx.pipeline(bl, "mm_hist"), bs = cx.pipeline(bl, "mm_scan"), bp = cx.pipeline(bl, "mm_place"),
                  bg = cx.pipeline(bl, "mm_gather_h"), bf = cx.pipeline(bl, "mm_gather");
    std::printf("device %s; mm_plan + mm_gather_h bitwise parity & benchmark\n", cx.device_name().c_str());
    const int TM = 16;
    const int nent = T * TOPK, nblk = (nent + 1023) / 1024, maxrows = nent + NE * TM;
    mtl::Buffer x = cx.alloc(4 * (size_t)T * n), xh = cx.alloc(2 * (size_t)T * n),
                ids = cx.alloc(4 * nent), w = cx.alloc(4 * nent),
                bcnt = cx.alloc(4 * (size_t)nblk * NE), boff = cx.alloc(4 * (size_t)nblk * NE), eoff = cx.alloc(4 * NE),
                tiles = cx.alloc(8 * (size_t)(maxrows / TM + 1)), meta = cx.alloc(64), stok = cx.alloc(4 * (size_t)maxrows),
                sw = cx.alloc(4 * (size_t)maxrows), pos = cx.alloc(4 * nent),
                XS = cx.alloc(2 * (size_t)maxrows * n), XS_ref = cx.alloc(2 * (size_t)maxrows * n);
    ref::Rng r(7);
    for (size_t i = 0; i < (size_t)T * n; ++i) {
        const float v = r.sym();
        x.as<float>()[i] = v;
        xh.as<uint16_t>()[i] = ref::f32_to_f16(v);
    }
    for (int t = 0; t < T; ++t) {
        std::set<int> used;
        float s = 0;
        for (int k = 0; k < TOPK; ++k) {
            int e;
            do e = (int)(r.next() % NE); while (used.count(e));
            used.insert(e);
            ids.as<int>()[t * TOPK + k] = e;
            s += (w.as<float>()[t * TOPK + k] = 0.2f + r.uni());
        }
        for (int k = 0; k < TOPK; ++k) w.as<float>()[t * TOPK + k] /= s;
    }
    MmPlanArgs pa{(uint32_t)nent, TOPK, (uint32_t)NE, (uint32_t)TM, (uint32_t)nblk};
    std::vector<mtl::Dispatch> ds_plan;
    mtl::Dispatch h; h.p = &phist; h.args = {mtl::Arg::B(ids), mtl::Arg::B(bcnt), mtl::Arg::V(pa)}; h.grid[0] = nblk; h.tg[0] = 1024; ds_plan.push_back(h);
    mtl::Dispatch sc; sc.p = &pscan; sc.args = {mtl::Arg::B(bcnt), mtl::Arg::B(boff), mtl::Arg::B(eoff), mtl::Arg::B(tiles), mtl::Arg::B(meta), mtl::Arg::V(pa), mtl::Arg::B(stok)}; sc.tg[0] = NE; ds_plan.push_back(sc);
    mtl::Dispatch pl; pl.p = &pplace; pl.args = {mtl::Arg::B(ids), mtl::Arg::B(w), mtl::Arg::B(boff), mtl::Arg::B(eoff), mtl::Arg::B(stok), mtl::Arg::B(sw), mtl::Arg::B(pos), mtl::Arg::V(pa)}; pl.grid[0] = nblk; pl.tg[0] = 1024; ds_plan.push_back(pl);
    bool ok_bits = true;
    if (baseline) {
        auto bd = ds_plan;
        bd[0].p = &bh; bd[1].p = &bs; bd[2].p = &bp;
        cx.run(bd);
        std::vector<mtl::Buffer*> outputs{&bcnt, &boff, &eoff, &tiles, &meta, &stok, &sw, &pos};
        std::vector<size_t> sizes{4u * nblk * NE, 4u * nblk * NE, 4u * NE,
            8u * (maxrows / TM + 1), 64, 4u * maxrows, 4u * maxrows, 4u * nent};
        std::vector<std::vector<uint8_t>> expected;
        for (size_t i = 0; i < outputs.size(); ++i) {
            const auto* bytes = outputs[i]->as<uint8_t>();
            expected.emplace_back(bytes, bytes + sizes[i]);
            std::memset(outputs[i]->ptr, 0, sizes[i]);
        }
        cx.run(ds_plan);
        const char* names[]{"bcnt", "boff", "eoff", "tiles", "meta", "srt_tok", "srt_w", "pos"};
        for (size_t i = 0; i < outputs.size(); ++i) {
            if (std::memcmp(outputs[i]->ptr, expected[i].data(), sizes[i]) != 0) {
                std::fprintf(stderr, "%s memcmp failed\n", names[i]);
                ok_bits = false;
            }
        }
        cx.warm(bd, 0.02);
        cx.warm(ds_plan, 0.02);
        std::vector<double> bt, ct;
        for (int i = 0; i < 7; ++i) {
            if (i & 1) {
                ct.push_back(cx.run(ds_plan, 8, false, true));
                bt.push_back(cx.run(bd, 8, false, true));
            } else {
                bt.push_back(cx.run(bd, 8, false, true));
                ct.push_back(cx.run(ds_plan, 8, false, true));
            }
        }
        std::printf("  plan baseline %.3f candidate %.3f us/L; all 8 buffers memcmp=%s\n",
                    median(bt) * 1e6 / 8, median(ct) * 1e6 / 8, ok_bits ? "PASS" : "FAIL");
    } else cx.run(ds_plan);
    const uint32_t tot_rows = meta.as<uint32_t>()[1], ne = n;
    const size_t xs_bytes = 2 * (size_t)tot_rows * n;
    bool ok_plan = true;
    std::vector<int> seen(tot_rows, 0);
    for (int i = 0; i < nent; ++i) {
        const int row = pos.as<int>()[i];
        if (row < 0 || row >= (int)tot_rows || stok.as<int>()[row] != i / TOPK) ok_plan = false;
        else seen[row]++;
    }
    for (uint32_t r0 = 0; r0 < tot_rows; ++r0)
        if (seen[r0] > 1 || (seen[r0] == 0 && stok.as<int>()[r0] != -1)) ok_plan = false;
    mtl::Dispatch d_ref; d_ref.p = &pgat_f; d_ref.args = {mtl::Arg::B(x), mtl::Arg::B(stok), mtl::Arg::B(XS_ref), mtl::Arg::V(ne)}; d_ref.grid[0] = tot_rows; d_ref.tg[0] = 256;
    cx.run({d_ref});
    auto b_ref = d_ref; b_ref.p = &bf; b_ref.args[2] = mtl::Arg::B(XS);
    cx.run({b_ref});
    bool ok_xs = std::memcmp(XS.ptr, XS_ref.ptr, xs_bytes) == 0;
    for (uint32_t tg_sz : {32u, 64u, 80u, 128u, 160u, 256u}) {
        std::memset(XS.ptr, 0x55, xs_bytes);
        mtl::Dispatch d; d.p = &pgat; d.args = {mtl::Arg::B(xh), mtl::Arg::B(stok), mtl::Arg::B(XS), mtl::Arg::V(ne)}; d.grid[0] = tot_rows; d.tg[0] = tg_sz;
        cx.run({d});
        if (std::memcmp(XS.ptr, XS_ref.ptr, xs_bytes) != 0) ok_xs = false;
        d.grid[0] = 1; d.indirect = &meta; d.ind_off = 8 * 4;
        cx.run({d});
        if (std::memcmp(XS.ptr, XS_ref.ptr, xs_bytes) != 0) ok_xs = false;
        auto b_gat = d; b_gat.p = &bg;
        cx.warm({b_gat, d}, 0.02);
        std::vector<double> bt, ct;
        for (int i = 0; i < 7; ++i) {
            if (i & 1) {
                ct.push_back(cx.run({d}, 8, false, true));
                bt.push_back(cx.run({b_gat}, 8, false, true));
            } else {
                bt.push_back(cx.run({b_gat}, 8, false, true));
                ct.push_back(cx.run({d}, 8, false, true));
            }
        }
        std::printf("  T=%d NE=%d rows=%u tg=%u baseline %.3f candidate %.3f us/L XS_memcmp=%s\n",
                    T, NE, tot_rows, tg_sz, median(bt) * 1e6 / 8, median(ct) * 1e6 / 8, ok_xs ? "PASS" : "FAIL");
    }
    auto combined = ds_plan;
    mtl::Dispatch gather = d_ref;
    gather.p = &pgat;
    gather.args[0] = mtl::Arg::B(xh);
    gather.args[2] = mtl::Arg::B(XS);
    gather.indirect = &meta; gather.ind_off = 8 * 4;
    combined.push_back(gather);
    std::vector<double> times;
    for (int i = 0; i < 3; ++i) times.push_back(cx.run(combined, 8, false, true));
    ok_xs &= std::memcmp(XS.ptr, XS_ref.ptr, xs_bytes) == 0;
    std::printf("  combined plan+gather tg256 %.3f us/L plan=%s XS_memcmp=%s\n",
                median(times) * 1e6 / 8, ok_plan ? "PASS" : "FAIL", ok_xs ? "PASS" : "FAIL");
    if (!ok_plan || !ok_xs || !ok_bits) return 1;
    return 0;
}

}  // namespace

int main(int argc, char** argv) {
    if (argc > 1 && std::string(argv[1]) == "--gather-bench") {
        return gather_bench(argc > 2 ? std::atoi(argv[2]) : 2048,
                            argc > 3 ? std::atoi(argv[3]) : 512,
                            argc > 4 ? std::atoi(argv[4]) : N_EMBD);
    }
    int gu = ref::Q2_0, dn = ref::Q2_0, T = 256, TM = 16, TN = 64, TK = 16, SG = 2, DTN = -1, DTK = -1, DSG = -1;
    bool bench = false, v2 = false, v3 = false, v4 = false, v5 = false;
    int MS = 1;
    bool gu_only = false, dn_only = false;
    const char* dump = nullptr, *ab_dir = nullptr;
    for (int i = 1; i < argc; ++i) {
        std::string a = argv[i];
        if (a == "--gu") gu = type_of(argv[++i]);
        else if (a == "--dn") dn = type_of(argv[++i]);
        else if (a == "--T" || a == "-T") T = std::atoi(argv[++i]);
        else if (a == "--tile") std::sscanf(argv[++i], "%d,%d,%d,%d", &TM, &TN, &TK, &SG);
        else if (a == "--dtile") {
            int m, n, k, s;
            const char* value = argv[++i];
            if (std::sscanf(value, "%d,%d,%d,%d", &m, &n, &k, &s) == 4) {
                if (m != TM) {
                    std::fprintf(stderr, "down tile TM must match gate/up TM\n");
                    return 2;
                }
                DTN = n; DTK = k; DSG = s;
            } else std::sscanf(value, "%d,%d,%d", &DTN, &DTK, &DSG);
        }
        else if (a == "--bench") bench = true;
        else if (a == "--gu-bench") bench = gu_only = true;   // --bench timing the gate/up GEMM alone
        else if (a == "--dn-bench") bench = dn_only = true;
        else if (a == "--v2") v2 = true;
        else if (a == "--v3") v2 = v3 = true;   // v2's gather + the register-operand GEMMs (TN per simdgroup)
        else if (a == "--v4") { v2 = v3 = v4 = true; MS = std::atoi(argv[++i]); }   // mm_gu4 with MS 16-row subtiles
        else if (a == "--v5") { v2 = v3 = true; v5 = true; }   // mm_gu5 (weight blocks staged in threadgroup memory)
        else if (a == "--dump") dump = argv[++i];   // HS and y to a file: bitwise comparison of two kernel versions
        else if (a == "--ab") ab_dir = argv[++i];
    }
    const bool isolated_bench = gu_only || dn_only;
    if (bench && T == 256) T = 2048;
    if (v4) TM = 16 * MS;   // the plan's tile height; mm_gu4's name carries MS instead
    mtl::Context cx;
    const int lib = cx.compile({"quant.metal", "units.metal", "mma.metal", "moe_mm.metal"}, (v3 && !v4 && !v5) ? "#define MM_V3_ONLY 1\n" : "");
    const int NE = bench ? 512 : 64, LAYERS = bench && !isolated_bench ? 4 : 1;
    const size_t gu_row = ref::row_bytes(gu, N_EMBD), d_row = ref::row_bytes(dn, N_FF);
    const size_t up_off = (size_t)N_FF * gu_row, down_off = 2 * up_off, blob = down_off + (size_t)N_EMBD * d_row;
    std::printf("device %s; %s/%s, T %d, %d experts, tile TM %d TN %d TK %d SG %d\n", cx.device_name().c_str(),
                ref::type_name(gu), ref::type_name(dn), T, NE, TM, TN, TK, SG);
    mtl::Buffer W = cx.alloc(blob * NE * LAYERS);
    {
        std::vector<uint8_t> tmp(blob * 16);
        for (int e = 0; e < 16; ++e) {
            ref::random_blocks(gu, tmp.data() + e * blob, 2 * N_FF, N_EMBD, 100 + e, 0.02f);
            ref::random_blocks(dn, tmp.data() + e * blob + down_off, N_EMBD, N_FF, 200 + e, 0.02f);
        }
        for (size_t e = 0; e < (size_t)NE * LAYERS; ++e) std::memcpy(W.as<uint8_t>() + e * blob, tmp.data() + (e % 16) * blob, blob);
    }
    const int nent = T * TOPK, nblk = (nent + 1023) / 1024;
    const int maxrows = nent + NE * TM;
    mtl::Buffer x = cx.alloc(4 * (size_t)T * N_EMBD), xh = cx.alloc(2 * (size_t)T * N_EMBD), ids = cx.alloc(4 * nent), w = cx.alloc(4 * nent),
                bcnt = cx.alloc(4 * (size_t)nblk * NE), boff = cx.alloc(4 * (size_t)nblk * NE), eoff = cx.alloc(4 * NE),
                tiles = cx.alloc(8 * (size_t)(maxrows / TM + 1)), meta = cx.alloc(64), stok = cx.alloc(4 * (size_t)maxrows),
                sw = cx.alloc(4 * (size_t)maxrows), pos = cx.alloc(4 * nent), H = cx.alloc(4 * (size_t)maxrows * N_FF),
                YE = cx.alloc(4 * (size_t)maxrows * N_EMBD), y = cx.alloc(4 * (size_t)T * N_EMBD),
                XS = cx.alloc(2 * (size_t)maxrows * N_EMBD), HS = cx.alloc(2 * (size_t)maxrows * N_FF);
    ref::Rng r(7);
    for (size_t i = 0; i < (size_t)T * N_EMBD; ++i) {
        const float v = r.sym();
        x.as<float>()[i] = v;
        xh.as<uint16_t>()[i] = ref::f32_to_f16(v);
    }
    for (int t = 0; t < T; ++t) {
        std::set<int> used;
        float s = 0;
        for (int k = 0; k < TOPK; ++k) {
            int e;
            do e = (int)(r.next() % NE); while (used.count(e));
            used.insert(e);
            ids.as<int>()[t * TOPK + k] = e;
            s += (w.as<float>()[t * TOPK + k] = 0.2f + r.uni());
        }
        for (int k = 0; k < TOPK; ++k) w.as<float>()[t * TOPK + k] /= s;
    }
    mtl::Pipeline phist = cx.pipeline(lib, "mm_hist"), pscan = cx.pipeline(lib, "mm_scan"), pplace = cx.pipeline(lib, "mm_place"),
                  pcomb = cx.pipeline(lib, "mm_combine");
    if (DTN < 0) { DTN = TN; DTK = TK; DSG = SG; }
    if (dn == ref::IQ4_NL && DTK < 32) DTK = 32;
    const std::string tl = std::to_string(TM) + "_" + std::to_string(TN) + "_" + std::to_string(TK) + "_" + std::to_string(SG);
    const std::string dtl = std::to_string(TM) + "_" + std::to_string(DTN) + "_" + std::to_string(DTK) + "_" + std::to_string(DSG);
    mtl::Constants probe;
    probe.i(5, std::getenv("MM_PROBE") ? std::atoi(std::getenv("MM_PROBE")) : 0);
    const bool k4 = coop_k4(cx);
    if (v3) std::printf("right-input layout: %s\n", k4 ? "runs of 4 k" : "other (per-element decode)");
    mtl::Pipeline pgu = v5 ? cx.pipeline(lib, "mm_gu5_" + std::to_string(gu) + "_" + std::to_string(TN) + "_" + std::to_string(SG),
                                         mtl::Constants().b(7, k4))
                      : v4 ? cx.pipeline(lib, "mm_gu4_" + std::to_string(gu) + "_" + std::to_string(MS) + tl.substr(tl.find('_')),
                                         mtl::Constants().b(7, k4).i(5, std::getenv("MM_PROBE") ? std::atoi(std::getenv("MM_PROBE")) : 0))
                      : v3 ? cx.pipeline(lib, "mm_gu3_" + std::to_string(gu) + "_" + tl, mtl::Constants().b(7, k4).i(5, std::getenv("MM_PROBE") ? std::atoi(std::getenv("MM_PROBE")) : 0).b(23, std::getenv("QUANT_CTAB") != nullptr))
                           : cx.pipeline(lib, (v2 ? "mm_gu2_" : "mm_gu_") + std::to_string(gu) + "_" + tl, probe),
                  pdn = v3 ? cx.pipeline(lib, "mm_down3_" + std::to_string(dn) + "_" + dtl, mtl::Constants().b(7, k4))
                           : cx.pipeline(lib, (v2 ? "mm_down2_" : "mm_down_") + std::to_string(dn) + "_" + (v2 ? dtl : tl)),
                  pgat = cx.pipeline(lib, "mm_gather_h");
    MmPlanArgs pa{(uint32_t)nent, TOPK, (uint32_t)NE, (uint32_t)TM, (uint32_t)nblk};
    auto encode = [&](int layer, std::vector<mtl::Dispatch>& ds, bool plan, bool gemm) {
        if (plan) {
            mtl::Dispatch h; h.p = &phist; h.args = {mtl::Arg::B(ids), mtl::Arg::B(bcnt), mtl::Arg::V(pa)}; h.grid[0] = nblk; h.tg[0] = 1024; ds.push_back(h);
            meta.as<uint32_t>()[2] = N_FF / TN;
            mtl::Dispatch sc; sc.p = &pscan; sc.args = {mtl::Arg::B(bcnt), mtl::Arg::B(boff), mtl::Arg::B(eoff), mtl::Arg::B(tiles), mtl::Arg::B(meta), mtl::Arg::V(pa), mtl::Arg::B(stok)}; sc.tg[0] = NE; ds.push_back(sc);
            mtl::Dispatch pl; pl.p = &pplace; pl.args = {mtl::Arg::B(ids), mtl::Arg::B(w), mtl::Arg::B(boff), mtl::Arg::B(eoff), mtl::Arg::B(stok), mtl::Arg::B(sw), mtl::Arg::B(pos), mtl::Arg::V(pa)}; pl.grid[0] = nblk; pl.tg[0] = 1024; ds.push_back(pl);
        }
        if (gemm) {
            MmArgs ma{(uint64_t)layer * NE * blob, (uint64_t)layer * NE * blob + up_off, (uint64_t)layer * NE * blob + down_off,
                      blob, blob, (uint32_t)gu_row, (uint32_t)d_row, N_EMBD, N_FF};
            if (v2) {
                const uint32_t ne = N_EMBD;
                mtl::Dispatch ga; ga.p = &pgat; ga.args = {mtl::Arg::B(xh), mtl::Arg::B(stok), mtl::Arg::B(XS), mtl::Arg::V(ne)};
                ga.grid[0] = meta.as<uint32_t>()[1]; ga.tg[0] = 256; ds.push_back(ga);
            }
            mtl::Dispatch g; g.p = &pgu;
            if (v2) g.args = {mtl::Arg::B(W), mtl::Arg::B(XS), mtl::Arg::B(HS), mtl::Arg::B(tiles), mtl::Arg::V(ma)};
            else g.args = {mtl::Arg::B(W), mtl::Arg::B(x), mtl::Arg::B(H), mtl::Arg::B(stok), mtl::Arg::B(tiles), mtl::Arg::V(ma)};
            g.grid[0] = N_FF / (v3 ? TN * SG : TN); g.grid[1] = meta.as<uint32_t>()[0]; g.tg[0] = 32 * SG; ds.push_back(g);
            mtl::Dispatch d; d.p = &pdn; d.args = {mtl::Arg::B(W), mtl::Arg::B(v2 ? HS : H), mtl::Arg::B(YE), mtl::Arg::B(tiles), mtl::Arg::V(ma)};
            d.grid[0] = N_EMBD / (v3 ? DTN * DSG : v2 ? DTN : TN); d.grid[1] = meta.as<uint32_t>()[0]; d.tg[0] = 32 * (v2 ? DSG : SG); ds.push_back(d);
            const uint32_t ca[4] = {(uint32_t)T, N_EMBD, TOPK, 0};
            mtl::Dispatch cb; cb.p = &pcomb; cb.args = {mtl::Arg::B(YE), mtl::Arg::B(pos), mtl::Arg::B(sw), mtl::Arg::B(y), mtl::Arg::V(ca)};
            cb.grid[0] = (T * N_EMBD + 255) / 256; cb.tg[0] = 256; ds.push_back(cb);
        }
    };
    {
        std::vector<mtl::Dispatch> ds;
        encode(0, ds, true, false);
        cx.run(ds);   // the plan first: the GEMMs' grid is the tile count
    }
    const uint32_t ntiles = meta.as<uint32_t>()[0];
    std::printf("plan: %u tiles of %d, %u padded rows for %d entries\n", ntiles, TM, meta.as<uint32_t>()[1], nent);
    auto compare_ab = [&](const std::vector<mtl::Dispatch>& ds, mtl::Pipeline& pgu_b, mtl::Pipeline& pdn_b, size_t pr) {
        std::vector<uint8_t> hs_a(HS.as<uint8_t>(), HS.as<uint8_t>() + 2 * pr * N_FF);
        std::vector<uint8_t> ye_a(YE.as<uint8_t>(), YE.as<uint8_t>() + 4 * pr * N_EMBD);
        std::vector<uint8_t> y_a(y.as<uint8_t>(), y.as<uint8_t>() + 4 * (size_t)T * N_EMBD);
        std::memset(HS.ptr, 0, hs_a.size());
        std::memset(YE.ptr, 0, ye_a.size());
        std::memset(y.ptr, 0, y_a.size());
        auto ds_b = ds;
        ds_b[v2 ? 1 : 0].p = &pgu_b;
        ds_b[v2 ? 2 : 1].p = &pdn_b;
        cx.run(ds_b);
        const bool ok_hs = std::memcmp(HS.ptr, hs_a.data(), hs_a.size()) == 0;
        const bool ok_ye = std::memcmp(YE.ptr, ye_a.data(), ye_a.size()) == 0;
        const bool ok_y = std::memcmp(y.ptr, y_a.data(), y_a.size()) == 0;
        std::printf("ab bitwise (T=%d): HS_memcmp=%s YE_memcmp=%s Y_memcmp=%s\n",
                T, ok_hs ? "PASS" : "FAIL", ok_ye ? "PASS" : "FAIL", ok_y ? "PASS" : "FAIL");
        return ok_hs && ok_ye && ok_y;
    };
    if (!bench) {
        // the plan: every entry placed once, rows of an expert contiguous and in entry order
        bool ok_plan = true;
        std::vector<int> seen(maxrows, 0);
        for (int i = 0; i < nent; ++i) {
            const int row = pos.as<int>()[i];
            if (row < 0 || row >= (int)meta.as<uint32_t>()[1] || stok.as<int>()[row] != i / TOPK) ok_plan = false;
            else seen[row]++;
        }
        for (int i = 0; i < maxrows; ++i) if (seen[i] > 1) ok_plan = false;
        std::vector<mtl::Dispatch> ds;
        encode(0, ds, false, true);
        cx.run(ds);
        const size_t pr = meta.as<uint32_t>()[1];
        if (dump) {
            FILE* f = std::fopen(dump, "wb");
            if (v2) std::fwrite(HS.as<uint8_t>(), 2, pr * N_FF, f);
            std::fwrite(YE.as<float>(), 4, pr * N_EMBD, f);
            std::fwrite(y.as<float>(), 4, (size_t)T * N_EMBD, f);
            std::fclose(f);
            return 0;
        }
        if (ab_dir) {
            const std::string kd(ab_dir);
            const int lib_b = cx.compile({kd + "/quant.metal", kd + "/units.metal", kd + "/mma.metal", kd + "/moe_mm.metal"}, (v3 && !v4 && !v5) ? "#define MM_V3_ONLY 1\n" : "");
            mtl::Pipeline pgu_b = cx.pipeline(lib_b, "mm_gu3_" + std::to_string(gu) + "_" + tl, mtl::Constants().b(7, k4).i(5, std::getenv("MM_PROBE") ? std::atoi(std::getenv("MM_PROBE")) : 0).b(23, std::getenv("QUANT_CTAB") != nullptr));
            mtl::Pipeline pdn_b = cx.pipeline(lib_b, "mm_down3_" + std::to_string(dn) + "_" + dtl, mtl::Constants().b(7, k4));
            return (ok_plan && compare_ab(ds, pgu_b, pdn_b, pr)) ? 0 : 1;
        }
        // reference: fp16 operands (x, weights, h), fp32+ accumulation; and the exact one
        double se16 = 0, se = 0, sr = 0;
        std::vector<float> wg((size_t)N_FF * N_EMBD), wu(wg.size()), wd((size_t)N_EMBD * N_FF);
        std::vector<double> yr16((size_t)T * N_EMBD, 0.0), yr((size_t)T * N_EMBD, 0.0);
        for (int e = 0; e < NE; ++e) {
            const uint8_t* b = W.as<uint8_t>() + e * blob;
            bool any = false;
            for (int i = 0; i < nent && !any; ++i) any = ids.as<int>()[i] == e;
            if (!any) continue;
            ref::dequant(gu, b, wg.data(), wg.size());
            ref::dequant(gu, b + up_off, wu.data(), wu.size());
            ref::dequant(dn, b + down_off, wd.data(), wd.size());
            for (int i = 0; i < nent; ++i) {
                if (ids.as<int>()[i] != e) continue;
                const int t = i / TOPK;
                const float* xt = x.as<float>() + (size_t)t * N_EMBD;
                std::vector<double> h16v(N_FF), hv(N_FF);
                for (int f = 0; f < N_FF; ++f) {
                    double g = 0, u = 0, g16 = 0, u16 = 0;
                    for (int k = 0; k < N_EMBD; ++k) {
                        const float xk = xt[k], xh = h16(xk);
                        g += (double)wg[(size_t)f * N_EMBD + k] * xk;
                        u += (double)wu[(size_t)f * N_EMBD + k] * xk;
                        g16 += (double)h16(wg[(size_t)f * N_EMBD + k]) * xh;
                        u16 += (double)h16(wu[(size_t)f * N_EMBD + k]) * xh;
                    }
                    hv[f] = g / (1 + std::exp(-g)) * u;
                    h16v[f] = h16((float)(g16 / (1 + std::exp(-g16)) * u16));
                }
                const double wt = w.as<float>()[i];
                for (int d = 0; d < N_EMBD; ++d) {
                    double s = 0, s16 = 0;
                    for (int f = 0; f < N_FF; ++f) {
                        s += (double)wd[(size_t)d * N_FF + f] * hv[f];
                        s16 += (double)h16(wd[(size_t)d * N_FF + f]) * h16v[f];
                    }
                    yr[(size_t)t * N_EMBD + d] += wt * s;
                    yr16[(size_t)t * N_EMBD + d] += wt * s16;
                }
            }
        }
        for (size_t i = 0; i < yr.size(); ++i) {
            const double g = y.as<float>()[i];
            se16 += (g - yr16[i]) * (g - yr16[i]);
            se += (g - yr[i]) * (g - yr[i]);
            sr += yr[i] * yr[i];
        }
        const double rel16 = std::sqrt(se16 / sr), rel = std::sqrt(se / sr);
        const bool ok = ok_plan && rel16 < 2e-4;
        std::printf("parity: plan %s; y vs fp16-operand reference rel L2 %.2e, vs exact float reference %.2e  %s\n",
                    ok_plan ? "ok" : "WRONG", rel16, rel, ok ? "PASS" : "FAIL");
        return ok ? 0 : 1;
    }
    // bench: the plan, the two GEMMs and the combine, LAYERS distinct weight sets rotated
    double flops = 0;
    {
        const double rows = meta.as<uint32_t>()[1];
        flops = 2.0 * nent * (2.0 * N_FF * N_EMBD + (double)N_EMBD * N_FF);
        std::printf("useful GFLOP per layer %.1f (padded rows %.0f = %.2fx entries)\n", flops / 1e9, rows, rows / nent);
    }
    auto timeit = [&](bool plan, bool gemm, int part) {
        std::vector<mtl::Dispatch> ds;
        const int repeats = isolated_bench ? 4 : 8;
        for (int l = 0; l < repeats; ++l) {
            std::vector<mtl::Dispatch> one;
            encode(l % LAYERS, one, plan, gemm);
            if (part >= 0 && gemm) { ds.push_back(one[part]); continue; }
            for (auto& d : one) ds.push_back(d);
        }
        cx.warm(ds, isolated_bench ? 0.1 : 0.3);
        std::vector<double> t;
        for (int i = 0; i < (isolated_bench ? 3 : 5); ++i) t.push_back(cx.run(ds) / repeats);
        return median(t);
    };
    const int o = v2 ? 1 : 0;
    if (isolated_bench) {
        std::vector<mtl::Dispatch> ds;
        encode(0, ds, false, true);
        cx.run(ds);
    }
    if (ab_dir && isolated_bench) {
        const std::string kd(ab_dir);
        const int lib_b = cx.compile({kd + "/quant.metal", kd + "/units.metal", kd + "/mma.metal", kd + "/moe_mm.metal"}, (v3 && !v4 && !v5) ? "#define MM_V3_ONLY 1\n" : "");
        mtl::Pipeline pgu_b = cx.pipeline(lib_b, "mm_gu3_" + std::to_string(gu) + "_" + tl, mtl::Constants().b(7, k4).i(5, std::getenv("MM_PROBE") ? std::atoi(std::getenv("MM_PROBE")) : 0).b(23, std::getenv("QUANT_CTAB") != nullptr));
        mtl::Pipeline pdn_b = cx.pipeline(lib_b, "mm_down3_" + std::to_string(dn) + "_" + dtl, mtl::Constants().b(7, k4));
        std::vector<mtl::Dispatch> parity;
        encode(0, parity, false, true);
        const size_t pr = meta.as<uint32_t>()[1];
        std::memset(HS.ptr, 0, 2 * pr * N_FF);
        std::memset(YE.ptr, 0, 4 * pr * N_EMBD);
        cx.run(parity);
        if (!compare_ab(parity, pgu_b, pdn_b, pr)) return 1;
        const int part = dn_only ? (o + 1) : o;
        std::vector<mtl::Dispatch> one;
        encode(0, one, false, true);
        std::vector<mtl::Dispatch> ds_a{one[part]};
        mtl::Dispatch db = one[part];
        db.p = dn_only ? &pdn_b : &pgu_b;
        std::vector<mtl::Dispatch> ds_b{db};
        cx.warm(ds_a, 0.05);
        cx.warm(ds_b, 0.05);
        std::vector<double> ta, tb;
        double best_pair = 1e30;
        for (int i = 0; i < 19; ++i) {
            double va, vb;
            if (i & 1) {
                vb = cx.run(ds_b, 2) / 2;
                va = cx.run(ds_a, 2) / 2;
            } else {
                va = cx.run(ds_a, 2) / 2;
                vb = cx.run(ds_b, 2) / 2;
            }
            best_pair = std::min(best_pair, std::min(va, vb));
            ta.push_back(va);
            tb.push_back(vb);
        }
        std::vector<double> fa, fb;
        for (size_t i = 0; i < ta.size(); ++i) {
            if (std::max(ta[i], tb[i]) <= 1.10 * best_pair) {
                fa.push_back(ta[i]);
                fb.push_back(tb[i]);
            }
        }
        if (fa.size() >= 3) { ta = fa; tb = fb; }
        const double ma = median(ta) * 1e6, mb = median(tb) * 1e6;
        const double mina = *std::min_element(ta.begin(), ta.end()) * 1e6;
        const double minb = *std::min_element(tb.begin(), tb.end()) * 1e6;
        std::printf("%s A/B (T=%d): A med=%.1f (min=%.1f) | B med=%.1f (min=%.1f) | diff med=%+.2f%% min=%+.2f%%\n",
                    dn_only ? "down" : "gate/up", T, ma, mina, mb, minb, (mb - ma) / ma * 100.0, (minb - mina) / mina * 100.0);
        return 0;
    }
    if (dn_only) {
        const double td = timeit(false, true, o + 1);
        std::printf("down %.1f us (%.2f TFLOP/s useful)\n", td * 1e6, 2.0 * nent * N_FF * N_EMBD / td / 1e12);
        return 0;
    }
    if (gu_only) {
        const double tg = timeit(false, true, o);
        std::printf("gate/up %.1f us (%.2f TFLOP/s useful)\n", tg * 1e6, 2.0 * nent * 2.0 * N_FF * N_EMBD / tg / 1e12);
        return 0;
    }
    const double tp = timeit(true, false, -1), tg = timeit(false, true, o), td = timeit(false, true, o + 1),
                 tc = timeit(false, true, o + 2), ta = timeit(false, true, -1);
    if (v2) std::printf("gather %.1f us\n", timeit(false, true, 0) * 1e6);
    const double fgu = 2.0 * nent * 2.0 * N_FF * N_EMBD, fdn = 2.0 * nent * N_EMBD * N_FF;
    std::printf("per layer: plan %.1f us, gate/up %.1f us (%.1f TFLOP/s useful), down %.1f us (%.1f TFLOP/s), combine %.1f us;"
                " GEMMs + combine %.1f us = %.1f TFLOP/s; %.1f us per token per layer -> 48 layers: %.0f tok/s (experts only)\n",
                tp * 1e6, tg * 1e6, fgu / tg / 1e12, td * 1e6, fdn / td / 1e12, tc * 1e6, ta * 1e6, flops / ta / 1e12,
                (ta + tp) / T * 1e6, T / (48 * (ta + tp)));
    return 0;
}