#include "qxmx_ref.h"

#include <algorithm>
#include <cmath>
#include <cstdio>
#include <cstring>

namespace qx {

/* ---------- f16 ---------- */
static float f16_to_f32(uint16_t bits) {
    uint32_t sign = (uint32_t)(bits & 0x8000) << 16;
    uint32_t exp  = (bits >> 10) & 0x1f, mant = bits & 0x3ff;
    union { uint32_t u; float f; } cv;
    if (exp == 0)  { cv.f = (float)mant * 0.00006103515625f; cv.u |= sign; return cv.f; }
    if (exp == 31) { cv.u = sign | 0x7f800000u | (mant ? 0x00400000u : 0); return cv.f; }
    cv.u = sign | ((exp + 112) << 23) | (mant << 13);
    return cv.f;
}

/* ---------- Q2 dequant ---------- */
/* GGUF Q2 2D tensor dims[0]=K (in, fast), dims[1]=M (out).
 * Supports both g64 (18 bytes/block) and g128 (34 bytes/block).
 * q in {0,1,2} -> (q-1) in {-1,0,+1}. */
void ref_dequant_row(const qx_weight* w, int m_idx, float* out) {
    int K = (int)w->dims[0];
    int qk2 = (w->nbytes * 64 / w->n_elem == 18) ? 64 : 128;
    int ngrp = K / qk2;
    if (qk2 == 64) {
        const qx_block_q2* row = (const qx_block_q2*)w->data + (size_t)m_idx * ngrp;
        for (int g = 0; g < ngrp; g++) {
            float d = f16_to_f32(row[g].d);
            for (int k = 0; k < 64; k++) {
                int q = (row[g].qs[k / 4] >> ((k % 4) * 2)) & 3;
                out[g * 64 + k] = d * (float)(q - 1);
            }
        }
    } else {
        const qx_block_q2_g128* row = (const qx_block_q2_g128*)w->data + (size_t)m_idx * ngrp;
        for (int g = 0; g < ngrp; g++) {
            float d = f16_to_f32(row[g].d);
            for (int k = 0; k < 128; k++) {
                int q = (row[g].qs[k / 4] >> ((k % 4) * 2)) & 3;
                out[g * 128 + k] = d * (float)(q - 1);
            }
        }
    }
}

/* W(M,N) @ x(N) -> y(M), W stored as GGUF Q2 [N][M] (dims[0]=N, dims[1]=M).
 * Dequant+dot per group, no intermediate.  Handles both g64 and g128. */
static void q2_matvec(const qx_weight* w, const float* x, float* y) {
    int N = (int)w->dims[0], M = (int)w->dims[1];
    int qk2 = (w->nbytes * 64 / w->n_elem == 18) ? 64 : 128;
    int ngrp = N / qk2;
    if (qk2 == 64) {
        const qx_block_q2* W = (const qx_block_q2*)w->data;
        for (int m = 0; m < M; m++) {
            float acc = 0.0f;
            const qx_block_q2* row = W + (size_t)m * ngrp;
            for (int g = 0; g < ngrp; g++) {
                float d = f16_to_f32(row[g].d);
                float p = 0.0f;
                for (int k = 0; k < 64; k++) {
                    int q = (row[g].qs[k / 4] >> ((k % 4) * 2)) & 3;
                    p += (float)(q - 1) * x[g * 64 + k];
                }
                acc += d * p;
            }
            y[m] = acc;
        }
    } else {
        const qx_block_q2_g128* W = (const qx_block_q2_g128*)w->data;
        for (int m = 0; m < M; m++) {
            float acc = 0.0f;
            const qx_block_q2_g128* row = W + (size_t)m * ngrp;
            for (int g = 0; g < ngrp; g++) {
                float d = f16_to_f32(row[g].d);
                float p = 0.0f;
                for (int k = 0; k < 128; k++) {
                    int q = (row[g].qs[k / 4] >> ((k % 4) * 2)) & 3;
                    p += (float)(q - 1) * x[g * 128 + k];
                }
                acc += d * p;
            }
            y[m] = acc;
        }
    }
}

/* F32 matvec: y(M) = W(M,N) @ x(N), W row-major [M][N]. */
/* F32 weight stored as GGUF tensor (1D or 2D). */

/* ---------- RMSNorm ----------
 * llama.cpp uses a single plain RMSNorm everywhere: out = w * norm(x).  The
 * GGUF converter bakes the +1 Gemma offset into the stored weight for every
 * norm EXCEPT the SSM post-delta norm (which is plain in HF too), so at runtime
 * all norms are identical.  Verified on the real GGUF: input/post/final/q/k
 * norms read ~1.0, ssm_norm ~0.87 (no offset baked). */
static void rmsnorm(const float* x, const float* w, float* out, int dim, float eps) {
    double ss = 0.0;
    for (int i = 0; i < dim; i++) ss += (double)x[i] * x[i];
    float inv = 1.0f / sqrtf((float)(ss / dim) + eps);
    for (int i = 0; i < dim; i++) out[i] = x[i] * inv * w[i];
}

/* ---------- activations ---------- */
static inline float silu(float x) { return x / (1.0f + expf(-x)); }
static inline float sigmoid(float x) { return 1.0f / (1.0f + expf(-x)); }

/* ---------- causal depthwise conv1d (Qwen3-Next / qwen35 convention) ----------
 * Per Aila + Candle bug #2: PyTorch Conv1d left-pads.  For kernel k=4, output at
 * step t = sum_{j=0..k-1} w[j] * x_hist[mod], where x_hist = [x[t-3], x[t-2],
 * x[t-1], x[t]].  GGUF conv1d.weight is stored channel-major [channels][k]
 * (dims[0]=k is the fast axis): w_conv[c*k+j], w[0]*oldest, w[k-1]*newest.
 * Conv history is [k-1][channels], most-recent-last. */
static void conv1d_step(const float* w_conv /*[channels][k], channel-major*/,
                        const float* x_hist /*[k-1][channels]*/, float* x_hist_new /*[k-1][channels]*/,
                        const float* x_new /*[channels]*/, float* out /*[channels]*/,
                        int k, int channels) {
    for (int c = 0; c < channels; c++) {
        float acc = 0.0f;
        /* history is x[t-(k-1)] .. x[t-1]; w[0] is oldest (channel-major weight) */
        for (int j = 0; j < k - 1; j++) acc += w_conv[c * k + j] * x_hist[j * channels + c];
        acc += w_conv[c * k + (k - 1)] * x_new[c];
        out[c] = acc;
    }
    /* shift history: drop oldest, append x_new */
    if (k > 1) {
        memmove(x_hist_new, x_hist + channels, sizeof(float) * channels * (k - 2));
        memcpy(x_hist_new + channels * (k - 2), x_new, sizeof(float) * channels);
    }
}

/* ---------- mRoPE with partial rotation (Candle bug #4) ----------
 * Apply RoPE only to the first rot_dim=64 of each head_dim=256 head, per the
 * mRoPE section widths [t,h,w,extra]=[11,11,10,0].  For text (1D positions)
 * the three section position-ids are all = pos.  Frequencies: inv_freq[i] =
 * 1 / (freq_base ^ (2i/rot_dim)), i in [0, rot_dim/2). */
static void apply_rope(float* q /*[n_qk_heads][head_dim]*/, float* k /*[n_kv_heads][head_dim]*/,
                       int n_qk, int n_kv, int head_dim, int rot_dim, int64_t pos,
                       double freq_base) {
    int half = rot_dim / 2;
    /* precompute (cos, sin) for the rot_dim/2 freqs at this position */
    std::vector<float> ctab(half), stab(half);
    for (int i = 0; i < half; i++) {
        double freq = 1.0 / pow(freq_base, (double)(2 * i) / rot_dim);
        double angle = pos * freq;
        ctab[i] = (float)cos(angle); stab[i] = (float)sin(angle);
    }
    auto rope_one = [&](float* v) {
        /* rotate pairs (v[i], v[i+half]) for i in [0,half) */
        for (int i = 0; i < half; i++) {
            float a = v[i], b = v[i + half];
            v[i]       = a * ctab[i] - b * stab[i];
            v[i+half]  = a * stab[i] + b * ctab[i];
        }
    };
    for (int h = 0; h < n_qk; h++) rope_one(q + h * head_dim);
    for (int h = 0; h < n_kv; h++) rope_one(k + h * head_dim);
}

/* ---------- Gated DeltaNet recurrence (decode, single token) ----------
 * Inputs: q,k [NKH][HD=128], v [NVH][HD=128], alpha/beta [NVH], already conv1d'd
 * and (for q,k) about to be L2-normalized inside.  State: [NVH][HD][HD].
 *
 * Math (Aila fused kernel, matches CKE REF):
 *   scale_qk = 1/sqrt(HD)
 *   q <- L2(q), k <- L2(k)         (per V-head's key group)
 *   decay = exp( ssm_a * softplus(alpha + dt_bias) )  // ssm_a = -exp(A_log), baked
 *   beta_s = sigmoid(beta)
 *   S <- decay * S                  (per V-head)
 *   kv_mem = S^T k
 *   delta  = (v - kv_mem) * beta_s
 *   S <- S + k outer delta
 *   out    = scale_qk * (S^T q)
 * The dt_bias here is ssm_dt.bias; alpha,beta are per-V-head projections of x.
 * Note: in this model ssm_alpha/ssm_beta are PROJECTIONS (out=48=NVH), not the
 * raw gate scalars -- they are applied per V-head after the in-proj split.
 * Exported (non-static) for the M1.1 chunked-scan equivalence harness. */
void deltanet_step(const float* q_in /*[NKH*HD]*/, const float* k_in,
                          const float* v_in /*[NVH*HD]*/, const float* alpha /*[NVH]*/,
                          const float* beta /*[NVH]*/, const float* a_log /*[NVH] = ssm_a*/,
                          const float* dt_bias /*[NVH]*/, float* state_io /*[NVH*HD*HD]*/,
                          float* out /*[NVH*HD]*/) {
    constexpr int HD = QX_SSM_STATE;      /* 128 */
    constexpr int NKH = QX_SSM_GROUPS;    /* 16  */
    constexpr int NVH = QX_SSM_DT_RANK;   /* 48  */
    const float scale_qk = 1.0f / sqrtf((float)HD);

    /* working q,k copied per V-head's key group (3 V-heads share a key head) */
    float q[NKH * HD], k[NKH * HD];
    memcpy(q, q_in, sizeof(q)); memcpy(k, k_in, sizeof(k));

    /* L2 normalize q,k per KEY head (shared across the 3 V-heads in its group) */
    for (int hk = 0; hk < NKH; hk++) {
        float* qp = q + hk * HD; float* kp = k + hk * HD;
        float qs = 0, ks = 0;
        for (int i = 0; i < HD; i++) { qs += qp[i]*qp[i]; ks += kp[i]*kp[i]; }
        float qinv = 1.0f / sqrtf(qs + QX_RMS_EPS);
        float kinv = 1.0f / sqrtf(ks + QX_RMS_EPS);
        for (int i = 0; i < HD; i++) { qp[i] *= qinv; kp[i] *= kinv; }
    }

    for (int hv = 0; hv < NVH; hv++) {
        int hk = hv % NKH;                   /* tiled V-head order: hv -> key head hv%NKH */
        const float* qp = q + hk * HD;
        const float* kp = k + hk * HD;
        const float* vp = v_in + hv * HD;
        float* S = state_io + (size_t)hv * HD * HD;
        float* outp = out + hv * HD;

        /* decay: HF g = -exp(A_log) * softplus(a + dt_bias).  The GGUF converter
         * bakes -exp(A_log) into ssm_a (conversion/qwen.py), so g = ssm_a * sp;
         * ssm_a is already negative.  decay = exp(g) in (0,1). */
        float apb = alpha[hv] + dt_bias[hv];
        float sp = apb > 20.0f ? apb : (apb < -20.0f ? expf(apb) : log1pf(expf(apb)));
        float g = a_log[hv] * sp;                   /* ssm_a = -exp(A_log), baked in */
        float decay = expf(g);
        float beta_s = sigmoid(beta[hv]);

        /* S <- decay*S; kv_mem = S^T k; delta = (v-kv_mem)*beta_s */
        float kv_mem[HD], delta[HD];
        for (int col = 0; col < HD; col++) {
            float acc = 0.0f;
            for (int row = 0; row < HD; row++) {
                float s = S[row * HD + col] * decay;
                S[row * HD + col] = s;
                acc += s * kp[row];
            }
            kv_mem[col] = acc;
            delta[col] = (vp[col] - acc) * beta_s;
        }
        /* S <- S + k outer delta; out = scale_qk * (S^T q) */
        for (int col = 0; col < HD; col++) {
            float acc = 0.0f;
            for (int row = 0; row < HD; row++) {
                float s = S[row * HD + col] + kp[row] * delta[col];
                S[row * HD + col] = s;
                acc += s * qp[row];
            }
            outp[col] = scale_qk * acc;
        }
    }
}

/* ---------- session ---------- */
void ref_session::init(const qx_model* m, int max_ctx) {
    int n_ssm = 0;
    for (int L = 0; L < QX_N_LAYER; L++) if (!m->blocks[L].is_attention) n_ssm++;
    int n_attn = QX_N_LAYER - n_ssm;
    ssm.assign(n_ssm, {});
    for (auto& s : ssm) s.resize(QX_SSM_DT_RANK, QX_SSM_STATE, QX_SSM_CONVD, QX_SSM_CONV_K);
    /* KV cache: [n_attn][max_ctx][KV_DIM], one slab per attention layer. */
    kv_k.assign((size_t)n_attn * max_ctx * QX_KV_DIM, 0.0f);
    kv_v.assign((size_t)n_attn * max_ctx * QX_KV_DIM, 0.0f);
    n_tokens = 0;
    this->max_ctx = max_ctx;
}

/* ---------- block forwards ---------- */
static void ffn_forward(const qx_ffn_block* f, const float* x, float* out) {
    std::vector<float> h(QX_D_MODEL), g(QX_FF_INTER), u(QX_FF_INTER), t(QX_FF_INTER);
    rmsnorm(x, (const float*)f->post_norm.data, h.data(), QX_D_MODEL, QX_RMS_EPS);
    q2_matvec(&f->gate, h.data(), g.data());
    q2_matvec(&f->up,   h.data(), u.data());
    for (int i = 0; i < QX_FF_INTER; i++) t[i] = silu(g[i]) * u[i];
    q2_matvec(&f->down, t.data(), out);
}

/* Attention block forward (decode, batch=1).  Updates KV cache for this layer
 * and this position.  kv_k/kv_v are [max_ctx][KV_DIM]; head_dim=256, GQA 24Q/4KV. */
static void attn_block_forward(const qx_attn_block* a, const float* x, float* out,
                               float* kv_k_layer, float* kv_v_layer, int64_t pos) {
    std::vector<float> h(QX_D_MODEL);
    rmsnorm(x, (const float*)a->attn_norm.data, h.data(), QX_D_MODEL, QX_RMS_EPS);

    /* GGUF attn_q [5120x12288] emits q AND the attention gate, laid out PER-HEAD
     * interleaved [h0_q(256)|h0_gate(256)|h1_q(256)|h1_gate(256)|...] (llama.cpp
     * ggml_view_3d with stride head_dim*2).  De-interleave into q[6144]/gate[6144]. */
    std::vector<float> qall(QX_Q_DIM * 2), q(QX_Q_DIM), gate(QX_Q_DIM);
    q2_matvec(&a->q, h.data(), qall.data());
    for (int hh = 0; hh < QX_N_HEAD; hh++) {
        memcpy(q.data()    + hh * QX_HEAD_DIM, qall.data() + hh * QX_HEAD_DIM * 2,
               sizeof(float) * QX_HEAD_DIM);
        memcpy(gate.data() + hh * QX_HEAD_DIM, qall.data() + hh * QX_HEAD_DIM * 2 + QX_HEAD_DIM,
               sizeof(float) * QX_HEAD_DIM);
    }
    std::vector<float> k(QX_KV_DIM), v(QX_KV_DIM);
    q2_matvec(&a->k, h.data(), k.data());
    q2_matvec(&a->v, h.data(), v.data());

    /* per-head q/k plain RMSNorm over head_dim=256 (q_norm/k_norm [256], +1 baked). */
    const float* qn = (const float*)a->q_norm.data;
    const float* kn = (const float*)a->k_norm.data;
    for (int hh = 0; hh < QX_N_HEAD; hh++)
        rmsnorm(q.data() + hh * QX_HEAD_DIM, qn, q.data() + hh * QX_HEAD_DIM, QX_HEAD_DIM, QX_RMS_EPS);
    for (int hh = 0; hh < QX_N_HEAD_KV; hh++)
        rmsnorm(k.data() + hh * QX_HEAD_DIM, kn, k.data() + hh * QX_HEAD_DIM, QX_HEAD_DIM, QX_RMS_EPS);

    /* Partial mRoPE (IMROPE, half-split pairs): rot_dim=64 only; text 1D pos. */
    apply_rope(q.data(), k.data(), QX_N_HEAD, QX_N_HEAD_KV, QX_HEAD_DIM, QX_ROT_DIM, pos, QX_ROPE_FREQ_BASE);

    /* append K,V to cache at this position */
    memcpy(kv_k_layer + pos * QX_KV_DIM, k.data(), sizeof(float) * QX_KV_DIM);
    memcpy(kv_v_layer + pos * QX_KV_DIM, v.data(), sizeof(float) * QX_KV_DIM);

    /* GQA attention: 24 Q heads, 4 KV heads (ratio 6).  head_dim 256. */
    std::vector<float> attn_out(QX_Q_DIM, 0.0f);
    float scale = 1.0f / sqrtf((float)QX_HEAD_DIM);
    for (int hh = 0; hh < QX_N_HEAD; hh++) {
        int hk = hh / (QX_N_HEAD / QX_N_HEAD_KV);
        const float* qph = q.data() + hh * QX_HEAD_DIM;
        std::vector<float> scores(pos + 1);
        float maxs = -1e30f;
        for (int64_t t = 0; t <= pos; t++) {
            const float* kt = kv_k_layer + t * QX_KV_DIM + hk * QX_HEAD_DIM;
            float s = 0.0f;
            for (int d = 0; d < QX_HEAD_DIM; d++) s += qph[d] * kt[d];
            s *= scale;
            scores[t] = s; if (s > maxs) maxs = s;
        }
        float sum = 0.0f;
        for (int64_t t = 0; t <= pos; t++) { scores[t] = expf(scores[t] - maxs); sum += scores[t]; }
        float inv = 1.0f / sum;
        std::vector<float> ao(QX_HEAD_DIM, 0.0f);
        for (int64_t t = 0; t <= pos; t++) {
            float w = scores[t] * inv;
            const float* vt = kv_v_layer + t * QX_KV_DIM + hk * QX_HEAD_DIM;
            for (int d = 0; d < QX_HEAD_DIM; d++) ao[d] += w * vt[d];
        }
        memcpy(attn_out.data() + hh * QX_HEAD_DIM, ao.data(), sizeof(float) * QX_HEAD_DIM);
    }

    /* output gate: sigmoid(gate) * attn_out, then o_proj */
    for (int i = 0; i < QX_Q_DIM; i++) attn_out[i] *= sigmoid(gate[i]);
    q2_matvec(&a->o, attn_out.data(), out);
}

/* SSM block forward (decode, batch=1). */
static void ssm_block_forward(const qx_ssm_block* s, ref_ssm_state* st,
                              const float* x, float* out) {
    std::vector<float> h(QX_D_MODEL), qkv(QX_SSM_CONVD), conv_out(QX_SSM_CONVD);
    rmsnorm(x, (const float*)s->attn_norm.data, h.data(), QX_D_MODEL, QX_RMS_EPS);

    /* fused qkv -> [q(2048) | k(2048) | v(6144)] = convd 10240 */
    q2_matvec(&s->qkv, h.data(), qkv.data());
    if (getenv("QXMX_DUMP_LAYERS")) {
        static bool ref_dbg = false;
        if (!ref_dbg) {
            ref_dbg = true;
            fprintf(stderr, "  REF qkv[0..2]: %.4f %.4f %.4f  qkv[2048..2050]: %.4f %.4f %.4f  qkv[4096..4098]: %.4f %.4f %.4f\n",
                    qkv[0],qkv[1],qkv[2],
                    qkv[QX_SSM_KEYD],qkv[QX_SSM_KEYD+1],qkv[QX_SSM_KEYD+2],
                    qkv[2*QX_SSM_KEYD],qkv[2*QX_SSM_KEYD+1],qkv[2*QX_SSM_KEYD+2]);
        }
    }

    /* causal conv1d on the full 10240-wide qkv (depthwise, k=4) */
    const float* wc = (const float*)s->conv1d.data;   /* [4][10240] */
    float hist_new[QX_SSM_CONV_K - 1][QX_SSM_CONVD];
    conv1d_step(wc, st->conv.data(), (float*)hist_new, qkv.data(), conv_out.data(),
                QX_SSM_CONV_K, QX_SSM_CONVD);
    memcpy(st->conv.data(), hist_new, sizeof(hist_new));

    /* SiLU on the conv output before splitting (Aila: silu_dev(v) on each convolved channel) */
    for (int i = 0; i < QX_SSM_CONVD; i++) conv_out[i] = silu(conv_out[i]);
    if (getenv("QXMX_DUMP_LAYERS")) {
        static bool conv_dbg = false;
        if (!conv_dbg) { conv_dbg = true;
            fprintf(stderr, "  REF post_conv[0..2]: %.4f %.4f %.4f  [2048..2050]: %.4f %.4f %.4f  [4096..4098]: %.4f %.4f %.4f\n",
                    conv_out[0],conv_out[1],conv_out[2],
                    conv_out[QX_SSM_KEYD],conv_out[QX_SSM_KEYD+1],conv_out[QX_SSM_KEYD+2],
                    conv_out[2*QX_SSM_KEYD],conv_out[2*QX_SSM_KEYD+1],conv_out[2*QX_SSM_KEYD+2]);
        }
    }

    const float* q = conv_out.data();
    const float* k = conv_out.data() + QX_SSM_KEYD;
    const float* v = conv_out.data() + QX_SSM_KEYD * 2;

    /* alpha, beta projections -> per-V-head scalars (NVH=48) */
    std::vector<float> alpha(QX_SSM_DT_RANK), beta(QX_SSM_DT_RANK);
    q2_matvec(&s->alpha, h.data(), alpha.data());   /* uses pre-conv h (Aila) */
    q2_matvec(&s->beta,  h.data(), beta.data());
    if (getenv("QXMX_DUMP_LAYERS")) {
        static bool ab_dbg = false;
        if (!ab_dbg) { ab_dbg = true;
            fprintf(stderr, "  REF alpha[0..3]: %.4f %.4f %.4f %.4f  beta[0..3]: %.4f %.4f %.4f %.4f\n",
                    alpha[0],alpha[1],alpha[2],alpha[3],
                    beta[0],beta[1],beta[2],beta[3]);
        }
    }

    /* gated delta-rule recurrence; out is [NVH*HD] = [6144] */
    std::vector<float> ssm_out(QX_SSM_VALD);
    deltanet_step(q, k, v, alpha.data(), beta.data(),
                  (const float*)s->a.data, (const float*)s->dt_bias.data,
                  st->state.data(), ssm_out.data());
    if (getenv("QXMX_DUMP_LAYERS")) {
        static bool dn_dbg = false;
        if (!dn_dbg) { dn_dbg = true;
            fprintf(stderr, "  REF ssm_out[0..2]: %.4f %.4f %.4f  [128..130]: %.4f %.4f %.4f  [6141..6143]: %.4f %.4f %.4f\n",
                    ssm_out[0],ssm_out[1],ssm_out[2],
                    ssm_out[128],ssm_out[129],ssm_out[130],
                    ssm_out[6141],ssm_out[6142],ssm_out[6143]);
        }
    }

    /* group RMSNorm over the state (ssm_norm is [128], applied per V-head's 128)
     * then output gate (sigmoid(attn_gate) * silu?), then ssm_out proj.  Per Aila:
     * the gate is silu(z) where z comes from the attn_gate projection of h. */
    const float* nw = (const float*)s->ssm_norm.data;  /* [128] */
    for (int hv = 0; hv < QX_SSM_DT_RANK; hv++) {
        rmsnorm(ssm_out.data() + hv * QX_SSM_STATE, nw, ssm_out.data() + hv * QX_SSM_STATE,
                QX_SSM_STATE, QX_RMS_EPS);
    }

    /* output gate z = attn_gate(h) -> [6144], gate = silu(z); out *= gate */
    std::vector<float> zg(QX_SSM_VALD);
    q2_matvec(&s->gate, h.data(), zg.data());
    if (getenv("QXMX_DUMP_LAYERS")) {
        static bool g_dbg = false;
        if (!g_dbg) { g_dbg = true;
            fprintf(stderr, "  REF zg[0..2]: %.4f %.4f %.4f  ssm_out_gated[0..2]: %.4f %.4f %.4f\n",
                    zg[0],zg[1],zg[2],
                    ssm_out[0]*silu(zg[0]), ssm_out[1]*silu(zg[1]), ssm_out[2]*silu(zg[2]));
        }
    }
    for (int i = 0; i < QX_SSM_VALD; i++) ssm_out[i] *= silu(zg[i]);

    q2_matvec(&s->out, ssm_out.data(), out);
}

/* ---------- top-level decode ---------- */
void ref_decode(const qx_model* m, ref_session* s, const float* x_in, float* logits_out) {
    std::vector<float> x(x_in, x_in + QX_D_MODEL), y(QX_D_MODEL);
    if (getenv("QXMX_DUMP_LAYERS")) {
        double norm = 0; for (int i = 0; i < QX_D_MODEL; i++) norm += (double)x[i]*x[i];
        norm = sqrt(norm / QX_D_MODEL);
        fprintf(stderr, "  EMB: x_rms=%.4f [%.4f %.4f %.4f]\n", norm, x[0], x[1], x[2]);
    }
    int64_t pos = s->n_tokens;
    int ssm_idx = 0, attn_idx = 0;
    size_t kv_stride = (size_t)s->max_ctx * QX_KV_DIM;
    std::vector<float> h1(QX_D_MODEL), f(QX_D_MODEL);
    for (int L = 0; L < QX_N_LAYER; L++) {
        const qx_block* b = &m->blocks[L];
        if (b->is_attention) {
            attn_block_forward(&b->attn, x.data(), y.data(),
                               s->kv_k.data() + (size_t)attn_idx * kv_stride,
                               s->kv_v.data() + (size_t)attn_idx * kv_stride, pos);
            attn_idx++;
        } else {
            ssm_block_forward(&b->ssm, &s->ssm[ssm_idx], x.data(), y.data());
            ssm_idx++;
        }
        /* residual after attention/SSM, then dense SwiGLU FFN with its own
         * post_norm + residual: out = (x + sublayer) + ffn(post_norm(x+sublayer)). */
        for (int i = 0; i < QX_D_MODEL; i++) h1[i] = x[i] + y[i];
        ffn_forward(&b->ffn, h1.data(), f.data());
        for (int i = 0; i < QX_D_MODEL; i++) x[i] = h1[i] + f[i];
        if (getenv("QXMX_DUMP_LAYERS")) {
            double norm = 0; for (int i = 0; i < QX_D_MODEL; i++) norm += (double)x[i]*x[i];
            norm = sqrt(norm / QX_D_MODEL);
            fprintf(stderr, "  L%02d: x_rms=%.4f  [%.4f %.4f %.4f]\n", L, norm, x[0], x[1], x[2]);
        }
    }
    /* output norm + lm_head */
    std::vector<float> h(QX_D_MODEL);
    rmsnorm(x.data(), (const float*)m->output_norm.data, h.data(), QX_D_MODEL, QX_RMS_EPS);
    q2_matvec(&m->output, h.data(), logits_out);
    s->n_tokens++;
}

}  // namespace qx