// test-delta-chunk: correctness harness for chunked delta-net (qwen4exp prefill).
//
// Compares three implementations of the fused Gated Delta Rule recurrence
// (see ggml_compute_forward_delta_net_f32 in ggml/src/ggml.c and
//  delta_net_recurrent_f32 in ggml/src/ggml-cuda/delta-net.cu):
//   (A) plain-C++ sequential reference (exact port, same op order),
//   (B) chunked reference (chunk size 64, boundary-state carry) — scaffolding
//       for the future parallel-scan CUDA kernel; must match (A) bitwise-ish,
//   (C) the repo's ggml op on the CPU backend (exercises the IQK AVX2 path
//       on Zen3) — must match (A) within FP reassociation tolerance.
//
// On a CUDA build the same binary gains --cuda mode (TODO): run (C) through
// the CUDA backend and compare vs (A), gating the chunked CUDA kernel.
//
// Usage: test-delta-chunk [--quick]   (quick trims the token sweep for CI)

#include "ggml.h"
#include "ggml-backend.h"

#include <cmath>
#include <cstdint>
#include <cstdio>
#include <cstring>
#include <string>
#include <vector>

// Set via DELTA_CHUNK env (default 64 = QWEN3NEXT_CHUNK_SIZE): sweeps chunk
// boundaries for the WY candidate without recompiling.
static int g_chunk = 0;
static bool g_use_cuda = false;

// ---------------------------------------------------------------- RNG ---
static uint64_t rng_state = 0x123456789abcdefULL;
static float rng_uniform(float lo, float hi) {
    rng_state ^= rng_state << 13;
    rng_state ^= rng_state >> 7;
    rng_state ^= rng_state << 17;
    const double u = (double)(rng_state >> 11) * (1.0 / 9007199254740992.0);
    return (float)(lo + u * (hi - lo));
}
static float rng_normal(void) {
    const float u1 = rng_uniform(1e-6f, 1.0f);
    const float u2 = rng_uniform(0.0f, 1.0f);
    return sqrtf(-2.0f * logf(u1)) * cosf(2.0f * (float)M_PI * u2);
}

// ---------------------------------------------------------------- case ---
struct Case {
    int hd;          // head dim (64 / 128)
    int nt;          // n_tokens
    int hk;          // n_k heads
    int hv;          // n_v heads
    int repeat;      // repeat_type 0 (divide) / 1 (mod)
    int nseq;        // n_seqs
    bool saved;      // allocate saved_steps
    bool hot_state;  // large init state (exercises +-1e6 clamp)
    bool hot_gate;   // large gate values (exercises exp(min(g,50)) cap)
    bool loose;      // explosive dynamics: relative-tolerance ggml check
};

struct Tensors {
    std::vector<float> q, k, v, g, beta, state;
};

// q,k: [hd,nt,hk,nseq]  v: [hd,nt,hv,nseq]  state: [hd,hd*hv,nseq]
// g, beta: canonical [nseq,nt,hv] order, idx=(b*nt+t)*hv+h — this matches the
// real graph (build_fused_delta_net permutes scalar gate [hv,t,s] and beta
// [hv,1,t,s] so element (t,h) sits at t*hv+h). The ggml tensor builder below
// transposes into ggml's own memory order; the reference uses canonical idx.
static void fill_tensors(const Case & c, Tensors & t) {
    const size_t nq = (size_t)c.hd * c.nt * c.hk * c.nseq;
    const size_t nv = (size_t)c.hd * c.nt * c.hv * c.nseq;
    const size_t ng = (size_t)c.nt * c.hv * c.nseq;
    const size_t ns = (size_t)c.hd * c.hd * c.hv * c.nseq;
    t.q.resize(nq); t.k.resize(nq); t.v.resize(nv);
    t.g.resize(ng); t.beta.resize(ng); t.state.resize(ns);
    for (auto & x : t.q) x = rng_normal();
    for (auto & x : t.k) x = rng_normal();
    // The real graph L2-normalizes q/k before the op (llama-delta-net.cpp),
    // and the IQK/CUDA kernels assume pre-normalized inputs — mirror that here
    // so the harness tests the kernels, not the (dead in practice) unnorm path.
    for (int b = 0; b < c.nseq; ++b) {
        for (int h = 0; h < c.hk; ++h) {
            for (int tt = 0; tt < c.nt; ++tt) {
                float * qq = t.q.data() + ((size_t)b * c.hk + h) * c.hd * c.nt + (size_t)tt * c.hd;
                float * kk = t.k.data() + ((size_t)b * c.hk + h) * c.hd * c.nt + (size_t)tt * c.hd;
                float qn = 0.0f, kn = 0.0f;
                for (int i = 0; i < c.hd; ++i) { qn += qq[i] * qq[i]; kn += kk[i] * kk[i]; }
                qn = 1.0f / sqrtf(qn + 1e-12f);
                kn = 1.0f / sqrtf(kn + 1e-12f);
                for (int i = 0; i < c.hd; ++i) { qq[i] *= qn; kk[i] *= kn; }
            }
        }
    }
    for (auto & x : t.v) x = rng_normal() * 0.5f;
    // Stable dynamics like the real model (forgetting gates: decay<1 almost
    // always). Explosive decay>1 makes ANY two orderings (even AVX2-FMA vs
    // scalar) diverge chaotically — proven by the production ggml op itself
    // failing tight comparison there — so torture gates stay in [-4,3]
    // (decay<=20: heavy dynamics, still value-meaningful). The exp-cap
    // fminf(g,50) path executes unconditionally; g>50 behavior in production
    // is clamp-rail agreement, covered by the hot_state cases below.
    for (auto & x : t.g) x = c.hot_gate ? rng_uniform(-4.0f, 3.0f) : rng_uniform(-3.0f, -0.1f);
    for (auto & x : t.beta) x = rng_uniform(-3.0f, 3.0f);
    // hot_state hits the +-1e6 rail deterministically on step 0 (init beyond
    // the rail) and then converges contractively (decay<1): rail firing is
    // covered, ulp flips at the boundary stay ~1e-7 relative and shrink.
    for (auto & x : t.state) x = c.hot_state ? rng_uniform(-1.5e6f, 1.5e6f) : rng_normal() * 0.5f;
}

// ------------------------------------------------------- sequential ref ---
// Exact port of ggml_compute_forward_delta_net_f32 (fused, scalar-gate path).
// Layouts mirror the kernel (NOT C-order): out is [b,t,h,r] with token stride
// hd*hv (graph views it as [S_v,H_v,nt,nseq]); saved steps are [t,b,h,state]
// with step stride hd*hd*hv*nseq; final state is [b,h,state].
struct Result {
    std::vector<float> out;    // [hd,nt,hv,nseq]
    std::vector<float> state;  // final [hd,hd*hv,nseq]
    std::vector<float> saved;  // [(nt-1)*hd*hd*hv*nseq] or empty
};

static inline float sigmoid_f(float x) { return 1.0f / (1.0f + expf(-x)); }

static void run_token(
        const Case & c, const Tensors & t, int b, int h, int tt,
        float * state /*[hd*hd]*/, float * out_t /*[hd]*/,
        const float scale) {
    const int hd = c.hd;
    // NB: repeat_type==1 maps head h to h % H_k (i.e. h % (n_heads/gqa_ratio)),
    // NOT h % (H_v/H_k) — matches kernel + ggml_compute_forward_delta_net_f32.
    const int hk = (c.repeat == 0) ? h / (c.hv / c.hk) : h % c.hk;
    const size_t qoff = ((size_t)b * c.hk + hk) * c.hd * c.nt + (size_t)tt * c.hd;
    const size_t voff = ((size_t)b * c.hv + h) * c.hd * c.nt + (size_t)tt * c.hd;
    const size_t goff = ((size_t)b * c.nt + tt) * c.hv + h;
    const float * q = t.q.data() + qoff;
    const float * k = t.k.data() + qoff;
    const float * v = t.v.data() + voff;
    const float beta = sigmoid_f(t.beta[goff]);
    const float decay = expf(fminf(t.g[goff], 50.0f));

    float qn = 0.0f, kn = 0.0f;
    for (int i = 0; i < hd; ++i) { qn += q[i] * q[i]; kn += k[i] * k[i]; }
    const float qni = 1.0f / sqrtf(qn + 1e-12f);
    const float kni = 1.0f / sqrtf(kn + 1e-12f);

    float score = 0.0f;
    for (int i = 0; i < hd; ++i) score += (k[i] * kni) * (q[i] * qni * scale);

    std::vector<float> vn(hd);
    for (int r = 0; r < hd; ++r) {
        float vp = 0.0f, ov = 0.0f;
        for (int col = 0; col < hd; ++col) {
            const float s = state[r + col * hd];
            vp += s * k[col];
            ov += s * q[col];
        }
        vn[r] = v[r] * beta - vp * beta * decay * kni;
        out_t[r] = ov * decay * qni * scale + vn[r] * score;
    }
    for (int col = 0; col < hd; ++col) {
        const float kc = k[col] * kni;
        for (int r = 0; r < hd; ++r) {
            float s = decay * state[r + col * hd] + vn[r] * kc;
            state[r + col * hd] = fminf(fmaxf(s, -1e6f), 1e6f);
        }
    }
}

static void run_sequential(const Case & c, const Tensors & t, Result & r) {
    const int hd = c.hd;
    const float scale = 1.0f / sqrtf((float)hd);
    const size_t stsz = (size_t)hd * hd;
    r.out.assign((size_t)hd * c.nt * c.hv * c.nseq, 0.0f);
    r.state.assign(stsz * c.hv * c.nseq, 0.0f);
    r.saved.clear();
    if (c.saved && c.nt > 1) r.saved.assign((size_t)(c.nt - 1) * stsz * c.hv * c.nseq, 0.0f);
    std::vector<float> st(stsz);
    for (int b = 0; b < c.nseq; ++b) {
        for (int h = 0; h < c.hv; ++h) {
            memcpy(st.data(), t.state.data() + ((size_t)b * c.hv + h) * stsz, stsz * sizeof(float));
            for (int tt = 0; tt < c.nt; ++tt) {
                float * out_t = r.out.data() + (((size_t)b * c.nt + tt) * c.hv + h) * hd;
                run_token(c, t, b, h, tt, st.data(), out_t, scale);
                if (c.saved && tt + 1 < c.nt) {
                    memcpy(r.saved.data() + (((size_t)tt * c.nseq + b) * c.hv + h) * stsz,
                           st.data(), stsz * sizeof(float));
                }
            }
            memcpy(r.state.data() + ((size_t)b * c.hv + h) * stsz, st.data(), stsz * sizeof(float));
        }
    }
}

// ---------------------------------------------------------- chunked ref ---
// Chunk size 64 (QWEN3NEXT_CHUNK_SIZE). Same per-token math, chunk-at-a-time
// scheduling with explicit boundary-state carry — the structure the future
// parallel-scan kernel must reproduce bit-compatibly.
#define DELTA_CHUNK 64

static void run_chunked(const Case & c, const Tensors & t, Result & r) {
    const int hd = c.hd;
    const float scale = 1.0f / sqrtf((float)hd);
    const size_t stsz = (size_t)hd * hd;
    r.out.assign((size_t)hd * c.nt * c.hv * c.nseq, 0.0f);
    r.state.assign(stsz * c.hv * c.nseq, 0.0f);
    r.saved.clear();
    if (c.saved && c.nt > 1) r.saved.assign((size_t)(c.nt - 1) * stsz * c.hv * c.nseq, 0.0f);
    std::vector<float> carry(stsz), st(stsz);
    for (int b = 0; b < c.nseq; ++b) {
        for (int h = 0; h < c.hv; ++h) {
            memcpy(carry.data(), t.state.data() + ((size_t)b * c.hv + h) * stsz, stsz * sizeof(float));
            for (int c0 = 0; c0 < c.nt; c0 += DELTA_CHUNK) {
                const int c1 = c0 + DELTA_CHUNK < c.nt ? c0 + DELTA_CHUNK : c.nt;
                memcpy(st.data(), carry.data(), stsz * sizeof(float)); // chunk-in = prev boundary
                for (int tt = c0; tt < c1; ++tt) {
                    float * out_t = r.out.data() + (((size_t)b * c.nt + tt) * c.hv + h) * hd;
                    run_token(c, t, b, h, tt, st.data(), out_t, scale);
                    if (c.saved && tt + 1 < c.nt) {
                        memcpy(r.saved.data() + (((size_t)tt * c.nseq + b) * c.hv + h) * stsz,
                               st.data(), stsz * sizeof(float));
                    }
                }
                memcpy(carry.data(), st.data(), stsz * sizeof(float)); // chunk-out boundary
            }
            memcpy(r.state.data() + ((size_t)b * c.hv + h) * stsz, carry.data(), stsz * sizeof(float));
        }
    }
}

// ------------------------------------------------------------- ggml op ---
// --cuda mode: run the SAME graph through the CUDA backend (registry lookup,
// no direct CUDA dependency — compiles everywhere) and compare vs the CPU
// result. This gates the chunked CUDA kernel: same op, same graph, the kernel
// dispatch inside is the only thing that differs. g/beta are built as
// permuted views exactly like build_fused_delta_net (never hand-patched
// strides, so backend tensor_set/alloc paths stay on supported ground).
static bool run_ggml_op_cuda(const Case & c, const Tensors & t, Result & r, std::string & err) {
    ggml_backend_t be = nullptr;
    for (size_t i = 0; i < ggml_backend_reg_get_count(); ++i) {
        const char * name = ggml_backend_reg_get_name(i);
        if (name && strstr(name, "CUDA")) {
            be = ggml_backend_reg_init_backend(i, nullptr);
            break;
        }
    }
    if (!be) { err = "no CUDA backend in registry (CPU-only build?)"; return false; }

    const int hd = c.hd;
    const size_t out_size = (size_t)hd * c.nt * c.hv * c.nseq;
    const size_t st_size = (size_t)hd * hd * c.hv * c.nseq;
    // no-alloc context: ggml_backend_alloc_ctx_tensors assigns the buffers.
    // (Pre-allocated tensors trip GGML_ASSERT(ggml_get_no_alloc(ctx)).)
    const size_t mem = (out_size + st_size) * 4
        + (t.q.size() + t.k.size() + t.v.size() + t.g.size() + t.beta.size() + t.state.size()) * 4
        + 64 * 1024 * 1024;
    struct ggml_init_params ip = { mem, nullptr, true };
    struct ggml_context * ctx = ggml_init(ip);
    if (!ctx) { err = "ggml_init failed"; ggml_backend_free(be); return false; }

    auto mkc = [&](const char * name, int64_t n0, int64_t n1, int64_t n2, int64_t n3) {
        struct ggml_tensor * ten = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, n0, n1, n2, n3);
        ggml_set_name(ten, name);
        return ten;
    };
    struct ggml_tensor * q = mkc("q", hd, c.nt, c.hk, c.nseq);
    struct ggml_tensor * k = mkc("k", hd, c.nt, c.hk, c.nseq);
    struct ggml_tensor * v = mkc("v", hd, c.nt, c.hv, c.nseq);
    // canonical [b,t,h] flat == contiguous base [hv,t,s] flat: permute to the
    // exact view topology production uses ([nt,1,hv,s] / [1,nt,hv,s]).
    struct ggml_tensor * gbase = mkc("gbase", c.hv, c.nt, c.nseq, 1);
    struct ggml_tensor * g = ggml_permute(ctx, gbase, 2, 0, 3, 1);
    ggml_set_name(g, "g");
    struct ggml_tensor * bbase = mkc("bbase", c.hv, 1, c.nt, c.nseq);
    struct ggml_tensor * beta = ggml_permute(ctx, bbase, 2, 0, 1, 3);
    ggml_set_name(beta, "beta");
    struct ggml_tensor * st = mkc("st", hd, hd * c.hv, 1, c.nseq);
    struct ggml_tensor * saved = nullptr;
    if (c.saved && c.nt > 1) {
        saved = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, (int64_t)(c.nt - 1) * st_size);
        ggml_set_name(saved, "saved");
    }
    struct ggml_tensor * res = ggml_delta_net(ctx, q, k, v, g, beta, st, saved);
    res->op_params[0] = c.repeat;

    if (!ggml_backend_supports_op(be, res)) { err = "CUDA backend declines GGML_OP_DELTA_NET"; goto fail; }
    {
        ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors(ctx, be);
        if (!buf) { err = "CUDA alloc failed"; goto fail; }
        ggml_backend_tensor_set(q, t.q.data(), 0, t.q.size() * 4);
        ggml_backend_tensor_set(k, t.k.data(), 0, t.k.size() * 4);
        ggml_backend_tensor_set(v, t.v.data(), 0, t.v.size() * 4);
        ggml_backend_tensor_set(gbase, t.g.data(), 0, t.g.size() * 4);
        ggml_backend_tensor_set(bbase, t.beta.data(), 0, t.beta.size() * 4);
        ggml_backend_tensor_set(st, t.state.data(), 0, t.state.size() * 4);
        struct ggml_cgraph * gf = ggml_new_graph_custom(ctx, 16, false);
        ggml_build_forward_expand(gf, res);
        ggml_backend_graph_compute(be, gf);
        r.out.assign(out_size, 0.0f);
        r.state.assign(st_size, 0.0f);
        ggml_backend_tensor_get(res, r.out.data(), 0, out_size * 4);
        // result tail holds the final state (no src7 here)
        {
            std::vector<float> tail(st_size);
            ggml_backend_tensor_get(res, tail.data(), out_size * 4, st_size * 4);
            r.state = std::move(tail);
        }
        // NOTE: ggml_backend_tensor_get on the whole 1-D result also covers it;
        // read the two regions explicitly to avoid stride assumptions.
        if (saved) {
            r.saved.assign((size_t)(c.nt - 1) * st_size, 0.0f);
            ggml_backend_tensor_get(saved, r.saved.data(), 0, r.saved.size() * 4);
        } else r.saved.clear();
        ggml_backend_buffer_free(buf);
    }
    ggml_backend_free(be);
    ggml_free(ctx);
    return true;
fail:
    ggml_backend_free(be);
    ggml_free(ctx);
    return false;
}

static bool run_ggml_op(const Case & c, const Tensors & t, Result & r, std::string & err) {
    if (g_use_cuda) return run_ggml_op_cuda(c, t, r, err);
    const int hd = c.hd;
    const size_t out_size = (size_t)hd * c.nt * c.hv * c.nseq;
    const size_t st_size = (size_t)hd * hd * c.hv * c.nseq;
    // result is 1-D [out + final state]; plus headroom ggml may need
    const size_t mem = (out_size + st_size + (c.saved && c.nt > 1 ? (size_t)(c.nt - 1) * st_size / (c.hv * c.nseq) * c.hv * c.nseq : 0)) * sizeof(float)
        + (t.q.size() + t.k.size() + t.v.size() + t.g.size() + t.beta.size() + t.state.size()) * sizeof(float)
        + 64 * 1024 * 1024;
    struct ggml_init_params ip = { mem, nullptr, false };
    struct ggml_context * ctx = ggml_init(ip);
    if (!ctx) { err = "ggml_init failed"; return false; }

    auto mk = [&](const char * name, ggml_type type, int64_t n0, int64_t n1, int64_t n2, int64_t n3, const void * data, size_t nbytes) {
        struct ggml_tensor * ten = ggml_new_tensor_4d(ctx, type, n0, n1, n2, n3);
        ggml_set_name(ten, name);
        memcpy(ten->data, data, nbytes);
        return ten;
    };
    struct ggml_tensor * q = mk("q", GGML_TYPE_F32, hd, c.nt, c.hk, c.nseq, t.q.data(), t.q.size() * 4);
    struct ggml_tensor * k = mk("k", GGML_TYPE_F32, hd, c.nt, c.hk, c.nseq, t.k.data(), t.k.size() * 4);
    struct ggml_tensor * v = mk("v", GGML_TYPE_F32, hd, c.nt, c.hv, c.nseq, t.v.data(), t.v.size() * 4);
    // g/beta: the IQK fast path (and the real graph's permuted views) address
    // element (t,h) at flat t*hv+h, NOT at the contiguous [nt,1,hv] position
    // t+h*nt. Keep canonical [b,t,h] data in place and patch strides to match:
    // g [nt,1,hv,nseq]: nb = [hv, nt*hv, 1, nt*hv] (floats).
    struct ggml_tensor * g = mk("g", GGML_TYPE_F32, c.nt, 1, c.hv, c.nseq, t.g.data(), t.g.size() * 4);
    g->nb[0] = (size_t)c.hv * 4;
    g->nb[1] = (size_t)c.nt * c.hv * 4;
    g->nb[2] = 4;
    g->nb[3] = (size_t)c.nt * c.hv * 4;
    struct ggml_tensor * beta = mk("beta", GGML_TYPE_F32, 1, c.nt, c.hv, c.nseq, t.beta.data(), t.beta.size() * 4);
    beta->nb[0] = 4;
    beta->nb[1] = (size_t)c.hv * 4;
    beta->nb[2] = 4;
    beta->nb[3] = (size_t)c.nt * c.hv * 4;
    struct ggml_tensor * st = mk("st", GGML_TYPE_F32, hd, hd * c.hv, 1, c.nseq, t.state.data(), t.state.size() * 4);
    struct ggml_tensor * saved = nullptr;
    std::vector<float> saved_buf;
    if (c.saved && c.nt > 1) {
        saved_buf.assign((size_t)(c.nt - 1) * st_size, 0.0f);
        saved = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, (int64_t)saved_buf.size());
        ggml_set_name(saved, "saved");
        memcpy(saved->data, saved_buf.data(), saved_buf.size() * 4);
    }
    struct ggml_tensor * res = ggml_delta_net(ctx, q, k, v, g, beta, st, saved);
    res->op_params[0] = c.repeat;

    struct ggml_cgraph * gf = ggml_new_graph_custom(ctx, 16, false);
    ggml_build_forward_expand(gf, res);
    if (ggml_graph_compute_with_ctx(ctx, gf, 1) != GGML_STATUS_SUCCESS) {
        err = "graph compute failed";
        ggml_free(ctx);
        return false;
    }
    r.out.assign((float *)res->data, (float *)res->data + out_size);
    r.state.assign((float *)res->data + out_size, (float *)res->data + out_size + st_size);
    if (saved) r.saved.assign((float *)saved->data, (float *)saved->data + saved_buf.size());
    else r.saved.clear();
    ggml_free(ctx);
    return true;
}

// ------------------------------------------------- candidate C: chunked WY ---
// True chunked formulation (triangular solve + GEMM assembly), the math the
// future parallel-scan CUDA kernel must reproduce. Per chunk [c0,c1) with
// incoming state S_in (0-based in-chunk indices, kh = L2-normalized k,
// qt = q*qni*scale exactly as run_token computes them):
//   P[i]    = prod_{m<=i} d_m
//   L[i][j] = b_i * P[i]/P[j] * (kh_i . kh_j), j<i
//   f_i     = b_i*v_i - b_i*P[i]*(S_in kh_i)
//   solve (I+L) e = f            (forward substitution)
//   B[i][j] = (P[i-1]/P[j]) * (kh_j . qt_i), j<i
//   a_i     = P[i-1]*(S_in qt_i) + sum_{j<i} B[i][j] e_j
//   o_i     = d_i*a_i + (kh_i . qt_i)*e_i
//   S_t     = d_t*S_{t-1} + e_t kh_t^T, clamped +-1e6 per step (same positions
//             as sequential, so clamp behavior is identical by construction)
// A second assembly path (single GEMM: S_C = P[C-1]*S_in + sum_i D[C-1,i] e_i
// kh_i^T) cross-checks the fast form the CUDA kernel uses when per-step
// checkpoints are not requested. sanc (below) records its max diff.

struct WYStats {
    int chunks = 0;
    int fast_chunks = 0; // took the triangular-solve path
    int fb_chunks = 0;   // fell back to sequential (wild dynamics guard)
    float gemm_assembly_diff = 0.0f; // max |stepwise S_C - GEMM S_C| seen
};

// Guard rails for the fast path (production rule for the CUDA kernel too):
// explosive dynamics (huge decay ratios or huge injections) make the
// unclamped intermediate formulation ill-conditioned vs the per-step-clamped
// sequential form. Beyond these bounds the chunk falls back to the exact
// sequential inner loop. Operating regime (|e|~O(1), ratios<=1 with decay<1)
// never trips it; torture dynamics do — which validates the fallback itself.
#define WY_GUARD_D 1e4f
#define WY_GUARD_E 1e4f

static void solve_chunk_wy(
        const Case & c, const Tensors & t, int b, int h,
        int c0, int c1, const float * s_in, float * s_out,
        float * out_base, float * saved_base, WYStats & st) {
    const int hd = c.hd;
    const int C = c1 - c0;
    const float scale = 1.0f / sqrtf((float)hd);
    // per-token ingredients (same factoring as run_token)
    // Absolute decay products are ALWAYS exp(logP[.]) (never progressive
    // products): bitwise identical to what the CUDA kernel stages in shared.
    std::vector<float> kh(C * hd), qt(C * hd), bv(C), dv(C), ck(C * hd), aq(C * hd);
    // log-decay prefix sums: every decay RATIO is evaluated as exp(logP[i]-logP[j]).
    // Raw products P[i] underflow to 0 over 64 steps with decay<1, turning later
    // P[i]/P[j] ratios into 0/0 = nan. exp-of-difference is exact-or-zero, which
    // is the mathematically right answer (true ratio ~0). Absolute P[i]
    // (=exp(logP[i])) may still underflow to 0 in f_i/S_C/output terms — also correct there.
    std::vector<float> logP(C);
    for (int i = 0; i < C; ++i) {
        const int tt = c0 + i;
        const int hk = (c.repeat == 0) ? h / (c.hv / c.hk) : h % c.hk;
        const float * q = t.q.data() + (((size_t)b * c.hk + hk) * c.nt + tt) * hd;
        const float * k = t.k.data() + (((size_t)b * c.hk + hk) * c.nt + tt) * hd;
        const float * v = t.v.data() + (((size_t)b * c.hv + h) * c.nt + tt) * hd;
        const float goff = t.g[((size_t)b * c.nt + tt) * c.hv + h];
        const float boff = t.beta[((size_t)b * c.nt + tt) * c.hv + h];
        float qn = 0.0f, kn = 0.0f;
        for (int d = 0; d < hd; ++d) { qn += q[d] * q[d]; kn += k[d] * k[d]; }
        qn = 1.0f / sqrtf(qn + 1e-12f);
        kn = 1.0f / sqrtf(kn + 1e-12f);
        bv[i] = 1.0f / (1.0f + expf(-boff));
        dv[i] = expf(fminf(goff, 50.0f));
        logP[i] = fminf(goff, 50.0f) + (i ? logP[i - 1] : 0.0f);
        for (int d = 0; d < hd; ++d) {
            kh[i * hd + d] = k[d] * kn;
            qt[i * hd + d] = q[d] * qn * scale;
        }
        // NB: ck/aq need kh/qt COMPLETE (separate loop — kh[e] for e>d is not
        // filled yet inside the loop above; folding them in silently zeroes
        // the tail of every dot product).
        for (int d = 0; d < hd; ++d) {
            float sk = 0.0f, sq = 0.0f;
            for (int e = 0; e < hd; ++e) { sk += s_in[d + e * hd] * kh[i * hd + e]; sq += s_in[d + e * hd] * qt[i * hd + e]; }
            ck[i * hd + d] = sk; // S_in kh_i
            aq[i * hd + d] = sq; // S_in qt_i
        }
    }
    // L[i][j] = b_i * D[i][j] * (kh_i . kh_j) and KQ[i][j] = kh_j . qt_i, j < i,
    // with D via exp-diff. Track the largest ratio: explosive dynamics make the
    // unclamped formulation ill-conditioned vs per-step clamping (see guard).
    std::vector<float> L(C * C, 0.0f), KQ(C * C, 0.0f);
    float maxD = 0.0f;
    for (int i = 0; i < C; ++i) {
        for (int j = 0; j < i; ++j) {
            float dot = 0.0f, kq = 0.0f;
            for (int d = 0; d < hd; ++d) { dot += kh[i * hd + d] * kh[j * hd + d]; kq += kh[j * hd + d] * qt[i * hd + d]; }
            const float Dij = expf(logP[i] - logP[j]);
            if (Dij > maxD) maxD = Dij;
            L[i * C + j] = bv[i] * Dij * dot;
            KQ[i * C + j] = kq;
        }
    }
    // f_i = b_i v_i - b_i P[i] ck_i ; solve (I+L) e = f
    std::vector<float> E(C * hd);
    float maxE = 0.0f;
    for (int i = 0; i < C; ++i) {
        const int tt = c0 + i;
        const float * v = t.v.data() + (((size_t)b * c.hv + h) * c.nt + tt) * hd;
        for (int d = 0; d < hd; ++d) {
            float e = bv[i] * v[d] - bv[i] * expf(logP[i]) * ck[i * hd + d];
            for (int j = 0; j < i; ++j) e -= L[i * C + j] * E[j * hd + d];
            E[i * hd + d] = e;
            const float ae = fabsf(e);
            if (ae > maxE) maxE = ae;
        }
    }
    const size_t stsz = (size_t)hd * hd;
    // Guard (production rule for the CUDA kernel too): wild dynamics fall back
    // to the exact sequential inner loop. Operating regime never trips it.
    if (maxD > WY_GUARD_D || maxE > WY_GUARD_E || !std::isfinite(maxD) || !std::isfinite(maxE)) {
        std::vector<float> S(stsz);
        memcpy(S.data(), s_in, stsz * sizeof(float));
        for (int i = 0; i < C; ++i) {
            const int tt = c0 + i;
            float * out_t = out_base + (((size_t)b * c.nt + tt) * c.hv + h) * hd;
            run_token(c, t, b, h, tt, S.data(), out_t, scale);
            if (saved_base && tt + 1 < c.nt) {
                memcpy(saved_base + (((size_t)tt * c.nseq + b) * c.hv + h) * stsz, S.data(), stsz * sizeof(float));
            }
        }
        memcpy(s_out, S.data(), stsz * sizeof(float));
        st.chunks++;
        st.fb_chunks++;
        return;
    }
    st.fast_chunks++;
    // outputs + stepwise state assembly (clamp positions identical to sequential)
    std::vector<float> S(stsz);
    memcpy(S.data(), s_in, stsz * sizeof(float));
    for (int i = 0; i < C; ++i) {
        const int tt = c0 + i;
        const float pim = i ? expf(logP[i - 1]) : 1.0f;
        const float logPim = i ? logP[i - 1] : 0.0f;
        float ci = 0.0f;
        for (int d = 0; d < hd; ++d) ci += kh[i * hd + d] * qt[i * hd + d];
        float * out_t = out_base + (((size_t)b * c.nt + tt) * c.hv + h) * hd;
        // a = P[i-1]*aq[i] + sum_{j<i} D[i-1][j] * KQ[i][j] * e_j (exp-diff)
        for (int d = 0; d < hd; ++d) {
            float a = pim * aq[i * hd + d];
            for (int j = 0; j < i; ++j) {
                a += expf(logPim - logP[j]) * KQ[i * C + j] * E[j * hd + d];
            }
            out_t[d] = dv[i] * a + ci * E[i * hd + d];
        }
        for (int col = 0; col < hd; ++col) {
            for (int row = 0; row < hd; ++row) {
                float s = dv[i] * S[row + col * hd] + E[i * hd + row] * kh[i * hd + col];
                S[row + col * hd] = fminf(fmaxf(s, -1e6f), 1e6f);
            }
        }
        if (saved_base && tt + 1 < c.nt) {
            memcpy(saved_base + (((size_t)tt * c.nseq + b) * c.hv + h) * stsz, S.data(), stsz * sizeof(float));
        }
    }
    memcpy(s_out, S.data(), stsz * sizeof(float));
    // cross-check: single-GEMM assembly of S_C (fast path for saved==NULL)
    {
        float md = 0.0f;
        for (int row = 0; row < hd; ++row) {
            for (int col = 0; col < hd; ++col) {
                float s = expf(logP[C - 1]) * s_in[row + col * hd];
                for (int i = 0; i < C; ++i) {
                    s += expf(logP[C - 1] - logP[i]) * E[i * hd + row] * kh[i * hd + col];
                }
                s = fminf(fmaxf(s, -1e6f), 1e6f);
                const float d = fabsf(s - S[row + col * hd]);
                if (d > md) md = d;
            }
        }
        if (md > st.gemm_assembly_diff) st.gemm_assembly_diff = md;
    }
    st.chunks++;
}

static void run_chunked_wy(const Case & c, const Tensors & t, Result & r, WYStats & st) {
    const int hd = c.hd;
    const int chunk = g_chunk > 0 ? g_chunk : DELTA_CHUNK;
    const size_t stsz = (size_t)hd * hd;
    r.out.assign((size_t)hd * c.nt * c.hv * c.nseq, 0.0f);
    r.state.assign(stsz * c.hv * c.nseq, 0.0f);
    r.saved.clear();
    if (c.saved && c.nt > 1) r.saved.assign((size_t)(c.nt - 1) * stsz * c.hv * c.nseq, 0.0f);
    std::vector<float> carry(stsz), snext(stsz);
    for (int b = 0; b < c.nseq; ++b) {
        for (int h = 0; h < c.hv; ++h) {
            memcpy(carry.data(), t.state.data() + ((size_t)b * c.hv + h) * stsz, stsz * sizeof(float));
            for (int c0 = 0; c0 < c.nt; c0 += chunk) {
                const int c1 = c0 + chunk < c.nt ? c0 + chunk : c.nt;
                solve_chunk_wy(c, t, b, h, c0, c1, carry.data(), snext.data(),
                               r.out.data(), c.saved ? r.saved.data() : nullptr, st);
                memcpy(carry.data(), snext.data(), stsz * sizeof(float));
            }
            memcpy(r.state.data() + ((size_t)b * c.hv + h) * stsz, carry.data(), stsz * sizeof(float));
        }
    }
}

// ---------------------------------------------------------------- check ---
static float max_abs_diff(const std::vector<float> & a, const std::vector<float> & b, size_t * at = nullptr) {
    float m = 0.0f;
    for (size_t i = 0; i < a.size(); ++i) {
        if (!std::isfinite(a[i]) || !std::isfinite(b[i])) { if (at) *at = i; return INFINITY; }
        const float d = fabsf(a[i] - b[i]);
        if (d > m) { m = d; if (at) *at = i; }
    }
    return m;
}

// Loose comparator for explosive-dynamics torture cases: same-sign inf counts
// as equal; otherwise relative tolerance on large values, absolute on small.
// Denominator floor is 1e3 (not 1): hot-state intermediates are ~1e5-1e6 scale
// where a single FMA-vs-mul+add ulp is ~1e-2 absolute, and near-cancellation
// can leave small residuals dominated by that noise — all expected FP noise,
// while genuine mapping bugs still show as O(1e3+) and fail loudly.
static float max_rel_diff(const std::vector<float> & a, const std::vector<float> & b, size_t * at = nullptr) {
    float m = 0.0f;
    for (size_t i = 0; i < a.size(); ++i) {
        const float x = a[i], y = b[i];
        if (!std::isfinite(x) || !std::isfinite(y)) {
            if (std::isfinite(x) != std::isfinite(y)) { if (at) *at = i; return INFINITY; }
            if ((x > 0) != (y > 0)) { if (at) *at = i; return INFINITY; }
            continue;
        }
        const float denom = fabsf(x) > 1e3f ? fabsf(x) : 1e3f;
        const float d = fabsf(x - y) / denom;
        if (d > m) { m = d; if (at) *at = i; }
    }
    return m;
}

static int failures = 0;
static void check_case(const Case & c, int idx) {
    Tensors t;
    fill_tensors(c, t);
    Result seq, chk, op, wy;
    WYStats wst;
    run_sequential(c, t, seq);
    run_chunked(c, t, chk);
    run_chunked_wy(c, t, wy, wst);
    const float d_out_cc = max_abs_diff(seq.out, chk.out);
    const float d_st_cc = max_abs_diff(seq.state, chk.state);
    float d_sv_cc = 0.0f;
    if (c.saved) d_sv_cc = max_abs_diff(seq.saved, chk.saved);
    // candidate C (WY): different summation order than sequential, so FP-level
    // tolerance even on stable cases; loose rules on torture cases
    const float d_out_wy = c.loose ? max_rel_diff(seq.out, wy.out) : max_abs_diff(seq.out, wy.out);
    const float d_st_wy = c.loose ? max_rel_diff(seq.state, wy.state) : max_abs_diff(seq.state, wy.state);
    float d_sv_wy = 0.0f;
    if (c.saved) d_sv_wy = c.loose ? max_rel_diff(seq.saved, wy.saved) : max_abs_diff(seq.saved, wy.saved);
    std::string err;
    if (!run_ggml_op(c, t, op, err)) {
        printf("case %2d hd=%d nt=%d hk=%d hv=%d rep=%d nseq=%d saved=%d hot=%d/%d  GGML-OP-FAIL: %s\n",
               idx, c.hd, c.nt, c.hk, c.hv, c.repeat, c.nseq, c.saved, c.hot_state, c.hot_gate, err.c_str());
        failures++;
        return;
    }
    const float d_out_op = c.loose ? max_rel_diff(seq.out, op.out) : max_abs_diff(seq.out, op.out);
    const float d_st_op = c.loose ? max_rel_diff(seq.state, op.state) : max_abs_diff(seq.state, op.state);
    float d_sv_op = 0.0f;
    size_t at_sv = 0;
    if (c.saved) d_sv_op = c.loose ? max_rel_diff(seq.saved, op.saved, &at_sv) : max_abs_diff(seq.saved, op.saved);

    // chunked must be exact (identical op order); ggml op allows FP reassociation.
    // Loose torture tolerance is 1e-2, not 1e-3: hot_gate chaos amplifies even
    // AVX2-FMA-vs-scalar 1-ulp diffs past 1e-3 (the production op itself does),
    // while genuine mapping bugs still read O(1e3+) — 5 orders of margin kept.
    // Candidate WY allows reorder noise (stable tol 1e-4); GEMM-assembly
    // cross-check must hold on all non-loose cases (fast path the CUDA kernel
    // uses when per-step checkpoints are off).
    const bool ok_cc = d_out_cc <= 1e-6f && d_st_cc <= 1e-6f && d_sv_cc <= 1e-6f;
    const float tol = c.loose ? 1e-2f : 2e-4f;
    const bool ok_op = d_out_op <= tol && d_st_op <= tol && d_sv_op <= tol;
    const float wytol = c.loose ? 1e-2f : 1e-4f;
    const bool ok_wy = d_out_wy <= wytol && d_st_wy <= wytol && d_sv_wy <= wytol
        && (c.loose || wst.gemm_assembly_diff <= 1e-3f)
        && (c.loose || wst.fb_chunks == 0); // stable suite must take fast path
    printf("case %2d hd=%d nt=%3d hk=%d hv=%d rep=%d nseq=%d saved=%d hot=%d/%d  "
           "chunked[out %.2e st %.2e sv %.2e] %s  ggml[out %.2e st %.2e sv %.2e] %s  "
           "wy[out %.2e st %.2e sv %.2e gemm %.2e fast %d/fb %d] %s\n",
           idx, c.hd, c.nt, c.hk, c.hv, c.repeat, c.nseq, c.saved, c.hot_state, c.hot_gate,
           d_out_cc, d_st_cc, d_sv_cc, ok_cc ? "OK " : "FAIL",
           d_out_op, d_st_op, d_sv_op, ok_op ? "OK " : "FAIL",
           d_out_wy, d_st_wy, d_sv_wy, wst.gemm_assembly_diff, wst.fast_chunks, wst.fb_chunks, ok_wy ? "OK " : "FAIL");
    if (!ok_cc || !ok_op || !ok_wy) failures++;
    if (c.loose && !ok_op) {
        const size_t stsz = (size_t)c.hd * c.hd;
        printf("      worst-sv idx %zu (t=%zu b=%zu h=%zu e=%zu): seq=%.6e ggml=%.6e\n",
               at_sv, at_sv / (stsz * c.nseq * c.hv), (at_sv / stsz / c.hv) % c.nseq,
               (at_sv / stsz) % c.hv, at_sv % stsz,
               seq.saved[at_sv], op.saved[at_sv]);
    }
}

int main(int argc, char ** argv) {
    bool quick = false;
    for (int i = 1; i < argc; ++i) {
        if (!strcmp(argv[i], "--quick")) quick = true;
        else if (!strcmp(argv[i], "--cuda")) g_use_cuda = true;
        else { printf("usage: %s [--quick] [--cuda]\n", argv[0]); return 1; }
    }
    if (const char * dc = getenv("DELTA_CHUNK")) {
        g_chunk = atoi(dc);
        if (g_chunk < 1 || g_chunk > 256) { printf("DELTA_CHUNK out of range 1..256\n"); return 1; }
        printf("DELTA_CHUNK=%d\n", g_chunk);
    }
    if (g_use_cuda) printf("--cuda: comparing CUDA backend vs sequential reference\n");
    std::vector<Case> cases;
    const int hds[] = {64, 128};
    const int toks_full[] = {1, 2, 7, 8, 9, 63, 64, 65, 100, 128, 200, 256};
    const int toks_quick[] = {1, 8, 64, 65, 200};
    const int * toks = quick ? toks_quick : toks_full;
    const int ntoks = quick ? 5 : 12;
    int idx = 0;
    for (int hd : hds) {
        for (int ti = 0; ti < ntoks; ++ti) {
            // (hk,hv): gqa 1 and 4, both repeat types; seqs 1..2; saved on/off
            cases.push_back({hd, toks[ti], 4, 4, 0, 1, toks[ti] > 1, false, false, false});
            if (ti % 3 == 0) cases.push_back({hd, toks[ti], 2, 8, 0, 1, toks[ti] > 1, false, false, false});
            if (ti % 3 == 1) cases.push_back({hd, toks[ti], 2, 8, 1, 2, toks[ti] > 1, false, false, false});
        }
    }
    // torture cases: clamp + exp-cap + saved steps (small nt: hot_gate dynamics
    // are explosive by construction, so these pin the paths, not precision)
    cases.push_back({128, 9, 2, 8, 0, 1, true, true, true, true});
    cases.push_back({64, 65, 4, 4, 1, 2, true, true, false, true});
    for (auto & c : cases) check_case(c, idx++);
    printf("%s: %d/%d cases passed\n", failures ? "FAIL" : "PASS", idx - failures, idx);
    return failures ? 1 : 0;
}