/* SPDX-License-Identifier: AGPL-3.0-or-later AND Apache-2.0 */
/* Copyright (C) 2026 David (Corey) Herrera */
/*
 * Portions derived from wkv7_cuda_fp32.cu in RWKV-LM
 * (https://github.com/BlinkDL/RWKV-LM), by Bo Peng (BlinkDL),
 * licensed under the Apache License, Version 2.0. Modified by
 * David (Corey) Herrera, 2026. See NOTICE and LICENSES/Apache-2.0.txt.
 */
/*
 * rx_pretrain.c — RWKV-7 pretraining kernels (WKV7 backward-through-time).
 *
 * Counterpart to rx_wkv7_one() in rwkvx.c.  Mirrors that forward's state
 * recurrence, [B][T][H][N] data layout, s_d[(h*N+i)*N+j] state convention and
 * per-head independence.  fp32 state throughout (train dtype).
 *
 * Attribution: the backward equations are a scalar-C transcription of BlinkDL's
 * Apache-2.0 RWKV-LM CUDA wkv7_cuda_fp32.cu backward_kernel, cross-checked
 * against the verified fp64 PyTorch oracle (wkv7_reference.py).  RWKV-LM is
 * Apache-2.0; direct transcription of the kernel math is permitted.
 *
 * See rx_pretrain.h and docs/wkv7_backward_reference.md for the §5 hand-offs.
 */
#include "rwkvx/rx_pretrain.h"

#include <stdlib.h>
#include <string.h>
#include <math.h>
#include <stdio.h>

/* Flat index helpers for the [B][T][H][N] input vectors and [B][H][N][N]
 * state.  Kept as static inlines so the layout lives in exactly one place. */
static inline int vx(int b, int t, int h, int n, int T, int H, int N) {
    return (((b * T + t) * H + h) * N + n);
}
static inline int sx(int b, int h, int i, int j, int H, int N) {
    return (((b * H + h) * N + i) * N + j);
}

/* ================================================================
 * Full-sequence forward (per head).  Zero-init or carry-in state.
 * ================================================================ */
int rx_wkv7_forward_seq(
    const float *r, const float *w, const float *k, const float *v,
    const float *a, const float *b, const float *s_init,
    float *y_out, float *s_final,
    int B, int T, int H, int N)
{
    float *S = (float *)malloc((size_t)B * H * N * N * sizeof(float));
    if (!S) return -1;
    if (s_init) memcpy(S, s_init, (size_t)B * H * N * N * sizeof(float));
    else        memset(S, 0, (size_t)B * H * N * N * sizeof(float));

    for (int bch = 0; bch < B; bch++) {
        for (int h = 0; h < H; h++) {
            for (int t = 0; t < T; t++) {
                for (int i = 0; i < N; i++) {
                    int sbase = sx(bch, h, i, 0, H, N);
                    /* sa[i] = sum_j a[j]*S[i][j]  (a = nkk) */
                    float sa = 0.0f;
                    for (int j = 0; j < N; j++)
                        sa += a[vx(bch, t, h, j, T, H, N)] * S[sbase + j];
                    float vi = v[vx(bch, t, h, i, T, H, N)];
                    float yy = 0.0f;
                    for (int j = 0; j < N; j++) {
                        int vj = vx(bch, t, h, j, T, H, N);
                        S[sbase + j] = S[sbase + j] * w[vj]
                                     + k[vj] * vi
                                     + sa * b[vj];
                        yy += S[sbase + j] * r[vj];
                    }
                    if (y_out) y_out[vx(bch, t, h, i, T, H, N)] = yy;
                }
            }
        }
    }
    if (s_final) memcpy(s_final, S, (size_t)B * H * N * N * sizeof(float));
    free(S);
    return 0;
}

/* ================================================================
 * Backward-through-time (reverse scan).
 *
 * Reverse scan (per (b,h), carrying dS = dL/dS_t, init 0):
 *   (1) dr[j]  = sum_i S_t[i][j]*dy[i];   dS[i][j] += dy[i]*r[j]
 *   (2) dk[j]  = sum_i dS[i][j]*v[i]
 *       dv[i]  = sum_j dS[i][j]*k[j]
 *       db[j]  = sum_i dS[i][j]*sa[i]
 *       dw[j]  = sum_i dS[i][j]*S_prev[i][j]      (EFFECTIVE decay; see §5.1)
 *       dsa[i] = sum_j dS[i][j]*b[j]
 *   (3) da[j]  = sum_i dsa[i]*S_prev[i][j]
 *       dS[i][j] = dS[i][j]*w[j] + dsa[i]*a[j]    (dS carry to t-1; §5.3)
 * After the loop dS is the grad w.r.t. the initial carry S_{-1}.
 * ================================================================ */
int rx_wkv7_backward(
    const float *r, const float *w, const float *k, const float *v,
    const float *a, const float *b, const float *dy, const float *s_init,
    float *dr, float *dw, float *dk, float *dv, float *da, float *db,
    float *dS_out,
    int B, int T, int H, int N)
{
    size_t state_sz = (size_t)B * H * N * N;
    /* Cache every pre-step state S_prev[t] and the per-step sa[t]. */
    float *S        = (float *)malloc(state_sz * sizeof(float));
    float *S_cache  = (float *)malloc((size_t)T * state_sz * sizeof(float));
    float *sa_cache = (float *)malloc((size_t)T * B * H * N * sizeof(float));
    float *dS       = (float *)calloc(state_sz, sizeof(float));
    float *St       = (float *)malloc((size_t)N * N * sizeof(float)); /* S_t block */
    float *dsa      = (float *)malloc((size_t)N * sizeof(float));
    if (!S || !S_cache || !sa_cache || !dS || !St || !dsa) {
        free(S); free(S_cache); free(sa_cache); free(dS); free(St); free(dsa);
        return -1;
    }

    /* ---- forward, caching S_prev and sa ---- */
    if (s_init) memcpy(S, s_init, state_sz * sizeof(float));
    else        memset(S, 0, state_sz * sizeof(float));
    for (int bch = 0; bch < B; bch++) {
        for (int h = 0; h < H; h++) {
            for (int t = 0; t < T; t++) {
                float *Scache_t = S_cache + (size_t)t * state_sz;
                float *sa_t     = sa_cache + (size_t)t * B * H * N;
                for (int i = 0; i < N; i++) {
                    int sbase = sx(bch, h, i, 0, H, N);
                    float sa = 0.0f;
                    for (int j = 0; j < N; j++) {
                        Scache_t[sbase + j] = S[sbase + j];     /* S_prev */
                        sa += a[vx(bch, t, h, j, T, H, N)] * S[sbase + j];
                    }
                    sa_t[(bch * H + h) * N + i] = sa;
                    float vi = v[vx(bch, t, h, i, T, H, N)];
                    for (int j = 0; j < N; j++) {
                        int vj = vx(bch, t, h, j, T, H, N);
                        S[sbase + j] = S[sbase + j] * w[vj]
                                     + k[vj] * vi
                                     + sa * b[vj];
                    }
                }
            }
        }
    }

    /* ---- reverse scan ---- */
    memset(dS, 0, state_sz * sizeof(float));
    for (int bch = 0; bch < B; bch++) {
        for (int h = 0; h < H; h++) {
            /* dS carry is per (b,h): reset its N*N block at each (b,h) start
             * by zeroing only that block (calloc already zeroed everything,
             * but each (b,h) is independent so re-zero for safety). */
            for (int i = 0; i < N; i++)
                for (int j = 0; j < N; j++)
                    dS[sx(bch, h, i, 0, H, N) + j] = 0.0f;

            for (int t = T - 1; t >= 0; t--) {
                const float *Scache_t = S_cache + (size_t)t * state_sz;
                const float *sa_t     = sa_cache + (size_t)t * B * H * N;

                /* recompute S_t block from S_prev, w, v, k, sa, b */
                for (int i = 0; i < N; i++) {
                    int sbase = sx(bch, h, i, 0, H, N);
                    float sa = sa_t[(bch * H + h) * N + i];
                    float vi = v[vx(bch, t, h, i, T, H, N)];
                    for (int j = 0; j < N; j++) {
                        int vj = vx(bch, t, h, j, T, H, N);
                        St[i * N + j] = Scache_t[sbase + j] * w[vj]
                                      + k[vj] * vi
                                      + sa * b[vj];
                    }
                }

                /* (1) read-out: dr and fold dy*r into dS */
                for (int j = 0; j < N; j++) {
                    int vj = vx(bch, t, h, j, T, H, N);
                    float drj = 0.0f;
                    for (int i = 0; i < N; i++)
                        drj += St[i * N + j] * dy[vx(bch, t, h, i, T, H, N)];
                    dr[vj] = drj;
                }
                for (int i = 0; i < N; i++) {
                    int sbase = sx(bch, h, i, 0, H, N);
                    float dyi = dy[vx(bch, t, h, i, T, H, N)];
                    for (int j = 0; j < N; j++)
                        dS[sbase + j] += dyi * r[vx(bch, t, h, j, T, H, N)];
                }

                /* (2) state-update grads (dk,dv,db,dw) and dsa */
                for (int j = 0; j < N; j++) {
                    int vj = vx(bch, t, h, j, T, H, N);
                    float dkj = 0.0f, dbj = 0.0f, dwj = 0.0f;
                    for (int i = 0; i < N; i++) {
                        float dSij = dS[sx(bch, h, i, 0, H, N) + j];
                        dkj += dSij * v[vx(bch, t, h, i, T, H, N)];
                        dbj += dSij * sa_t[(bch * H + h) * N + i];
                        dwj += dSij * Scache_t[sx(bch, h, i, 0, H, N) + j];
                    }
                    dk[vj] = dkj;
                    db[vj] = dbj;
                    dw[vj] = dwj;   /* grad w.r.t. EFFECTIVE decay (bare) */
                }
                for (int i = 0; i < N; i++) {
                    int sbase = sx(bch, h, i, 0, H, N);
                    float dvi = 0.0f, dsai = 0.0f;
                    for (int j = 0; j < N; j++) {
                        int vj = vx(bch, t, h, j, T, H, N);
                        dvi  += dS[sbase + j] * k[vj];
                        dsai += dS[sbase + j] * b[vj];
                    }
                    dv[vx(bch, t, h, i, T, H, N)] = dvi;
                    dsa[i] = dsai;
                }

                /* (3) in-context term: da, then dS carry -> t-1 */
                for (int j = 0; j < N; j++) {
                    int vj = vx(bch, t, h, j, T, H, N);
                    float daj = 0.0f;
                    for (int i = 0; i < N; i++)
                        daj += dsa[i] * Scache_t[sx(bch, h, i, 0, H, N) + j];
                    da[vj] = daj;
                }
                for (int i = 0; i < N; i++) {
                    int sbase = sx(bch, h, i, 0, H, N);
                    for (int j = 0; j < N; j++) {
                        int vj = vx(bch, t, h, j, T, H, N);
                        dS[sbase + j] = dS[sbase + j] * w[vj]
                                      + dsa[i] * a[vj];
                    }
                }
            }
        }
    }

    if (dS_out) memcpy(dS_out, dS, state_sz * sizeof(float));

    free(S); free(S_cache); free(sa_cache); free(dS); free(St); free(dsa);
    return 0;
}

/* ================================================================
 * RWKV-7 TIME-MIX BLOCK forward/backward.
 *
 * Mirrors rwkvx.c `rwkv7_time_mix_fast` (scalar #else path) exactly; the
 * recurrence is delegated to rx_wkv7_forward_seq / rx_wkv7_backward above.
 * Everything is fp32 in the compute path.  Weights are row-major [in,out]
 * exactly as fast_matvec consumes them.
 * ================================================================ */

#define RX_WDECAY_COEFF 0.606531f   /* == exp(-0.5), foundry's decay scale */

/* Per-token stack-scratch bounds for the backward.  C (=n_head*head_size) and
 * head_size must not exceed these in the gradient-check harness. */
#define RX_TM_MAXC 4096
#define RX_TM_MAXN 4096

/* out[j] = sum_i x[i]*W[i*cols+j]   (matches fast_matvec) */
static void mv_fwd(const float *x, const float *W, float *out,
                   int rows, int cols) {
    for (int j = 0; j < cols; j++) out[j] = 0.0f;
    for (int i = 0; i < rows; i++) {
        float xi = x[i];
        const float *Wi = &W[(size_t)i * cols];
        for (int j = 0; j < cols; j++) out[j] += xi * Wi[j];
    }
}

/* Backward of mv_fwd.  Accumulates: dx[i] += sum_j dout[j]*W[i*cols+j];
 * dW[i*cols+j] += x[i]*dout[j]  (dW may be NULL to skip). */
static void mv_bwd(const float *x, const float *W, const float *dout,
                   float *dx, float *dW, int rows, int cols) {
    for (int i = 0; i < rows; i++) {
        float acc = 0.0f;
        const float *Wi = &W[(size_t)i * cols];
        float xi = x[i];
        if (dW) {
            float *dWi = &dW[(size_t)i * cols];
            for (int j = 0; j < cols; j++) {
                acc += dout[j] * Wi[j];
                dWi[j] += xi * dout[j];
            }
        } else {
            for (int j = 0; j < cols; j++) acc += dout[j] * Wi[j];
        }
        dx[i] += acc;
    }
}

static inline float sigmoidf_(float z) { return 1.0f / (1.0f + expf(-z)); }

/* L2 normalize per head (scalar mirror of fast_l2_normalize_per_head). */
static void l2norm_head(const float *x, float *out, int H, int N) {
    for (int h = 0; h < H; h++) {
        const float *hx = &x[h * N];
        float *ho = &out[h * N];
        float ss = 0.0f;
        for (int i = 0; i < N; i++) ss += hx[i] * hx[i];
        float inv = 1.0f / (sqrtf(ss) + 1e-12f);
        for (int i = 0; i < N; i++) ho[i] = hx[i] * inv;
    }
}

/* Build the six WKV7 inputs (R,W,K,V,A,B) + fill optional caches for one token.
 * Returns nothing; writes into the *_o output slices (each length c) and the
 * per-token cache slices when non-NULL. */
typedef struct {
    /* caches (each length c unless noted) */
    float *shr,*shk,*shv,*shw,*sha,*shg;
    float *rvec,*kvec,*vvec;
    float *sigw;             /* sigmoid pre-decay */
    float *avec, *gvec;
    float *tk, *kk;
    float *lw, *tw;          /* [Dw] */
    float *la;               /* [Da] */
    float *lg, *sg;          /* [Dg] */
    float *lv, *blend;       /* [Dv]/[c], layer>0 */
} tm_cache;

/* Preprocessing forward for one token.  xx/prev are [c]; v_first_tok is the
 * per-token v_first slice (read for layer>0, written for layer 0). */
static void tm_pre_forward(const rx_tm_params *P, const float *xx,
    const float *prev, float *v_first_tok,
    float *R, float *W, float *K, float *V, float *A, float *B,
    const tm_cache *c)
{
    int C = P->c, H = P->n_head, N = P->head_size;
    /* token shifts */
    for (int i = 0; i < C; i++) {
        float d = prev[i] - xx[i];
        c->shr[i] = xx[i] + d * P->x_r[i];
        c->shk[i] = xx[i] + d * P->x_k[i];
        c->shv[i] = xx[i] + d * P->x_v[i];
        c->shw[i] = xx[i] + d * P->x_w[i];
        c->sha[i] = xx[i] + d * P->x_a[i];
        c->shg[i] = xx[i] + d * P->x_g[i];
    }
    /* R/K/V projections */
    mv_fwd(c->shr, P->receptance_w, c->rvec, C, C);
    mv_fwd(c->shk, P->key_w,        c->kvec, C, C);
    mv_fwd(c->shv, P->value_w,      c->vvec, C, C);
    for (int i = 0; i < C; i++) R[i] = c->rvec[i];

    /* decay: w = exp(-coeff * sigmoid(w0 + tanh(shw@w1)@w2)) */
    mv_fwd(c->shw, P->w1, c->lw, C, P->Dw);
    for (int d = 0; d < P->Dw; d++) c->tw[d] = tanhf(c->lw[d]);
    mv_fwd(c->tw, P->w2, W, P->Dw, C);           /* ow -> W temporarily */
    for (int i = 0; i < C; i++) {
        float sig = sigmoidf_(P->w0[i] + W[i]);
        c->sigw[i] = sig;
        W[i] = expf(-RX_WDECAY_COEFF * sig);
    }
    /* gate a: sigmoid(a0 + (sha@a1)@a2) */
    mv_fwd(c->sha, P->a1, c->la, C, P->Da);
    mv_fwd(c->la, P->a2, c->avec, P->Da, C);     /* oa -> avec temporarily */
    for (int i = 0; i < C; i++) c->avec[i] = sigmoidf_(P->a0[i] + c->avec[i]);

    /* gate g: (sigmoid(shg@g1))@g2 */
    mv_fwd(c->shg, P->g1, c->lg, C, P->Dg);
    for (int d = 0; d < P->Dg; d++) c->sg[d] = sigmoidf_(c->lg[d]);
    mv_fwd(c->sg, P->g2, c->gvec, P->Dg, C);

    /* key norm: kk = L2norm_per_head(k * k_k) */
    for (int i = 0; i < C; i++) c->tk[i] = c->kvec[i] * P->k_k[i];
    l2norm_head(c->tk, c->kk, H, N);

    /* key mod: K = k * (1 + (a-1)*k_a) */
    for (int i = 0; i < C; i++)
        K[i] = c->kvec[i] * (1.0f + (c->avec[i] - 1.0f) * P->k_a[i]);

    /* value residual gating */
    if (P->layer_id == 0) {
        for (int i = 0; i < C; i++) { V[i] = c->vvec[i]; v_first_tok[i] = c->vvec[i]; }
    } else {
        mv_fwd(c->shv, P->v1, c->lv, C, P->Dv);
        mv_fwd(c->lv, P->v2, c->blend, P->Dv, C); /* ov -> blend temporarily */
        for (int i = 0; i < C; i++) {
            float bl = sigmoidf_(P->v0[i] + c->blend[i]);
            c->blend[i] = bl;
            V[i] = c->vvec[i] + (v_first_tok[i] - c->vvec[i]) * bl;
        }
    }
    /* WKV7 composites */
    for (int i = 0; i < C; i++) { A[i] = -c->kk[i]; B[i] = c->kk[i] * c->avec[i]; }
}

/* group-norm forward for one token; stores per-head mean/rstd. */
static void gn_forward(const float *y, const float *w, const float *b,
                       float *out, float *mean_o, float *rstd_o, int H, int N) {
    for (int h = 0; h < H; h++) {
        const float *gy = &y[h * N];
        float *go = &out[h * N];
        float mean = 0.0f;
        for (int i = 0; i < N; i++) mean += gy[i];
        mean /= N;
        float var = 0.0f;
        for (int i = 0; i < N; i++) { float d = gy[i] - mean; var += d * d; }
        var /= N;
        float rstd = 1.0f / sqrtf(var + 64e-5f);
        mean_o[h] = mean; rstd_o[h] = rstd;
        for (int i = 0; i < N; i++)
            go[i] = (gy[i] - mean) * rstd * w[h * N + i] + b[h * N + i];
    }
}

int rx_timemix_forward_seq(
    const rx_tm_params *P, const float *x, const float *x_prev_init,
    float *v_first, float *out, int B, int T)
{
    int C = P->c, H = P->n_head, N = P->head_size;
    size_t BT = (size_t)B * T, tot = BT * C;
    int maxD = P->Dw; if (P->Da>maxD) maxD=P->Da; if (P->Dg>maxD) maxD=P->Dg;
    if (P->layer_id>0 && P->Dv>maxD) maxD=P->Dv;

    float *R=malloc(tot*sizeof(float)), *W=malloc(tot*sizeof(float));
    float *K=malloc(tot*sizeof(float)), *V=malloc(tot*sizeof(float));
    float *A=malloc(tot*sizeof(float)), *Bb=malloc(tot*sizeof(float));
    float *y=malloc(tot*sizeof(float)), *gv=malloc(tot*sizeof(float));
    /* single-token scratch for all caches except gvec (stored full in gv) */
    float *cb=calloc((size_t)(15*C + 6*(size_t)maxD), sizeof(float));
    float *ygn=malloc((size_t)C*sizeof(float));
    float *mh=malloc((size_t)H*sizeof(float)), *rh=malloc((size_t)H*sizeof(float));
    if(!R||!W||!K||!V||!A||!Bb||!y||!gv||!cb||!ygn||!mh||!rh){
        free(R);free(W);free(K);free(V);free(A);free(Bb);free(y);free(gv);
        free(cb);free(ygn);free(mh);free(rh);return -1;
    }
    float *p=cb; tm_cache c;
    c.shr=p;p+=C;c.shk=p;p+=C;c.shv=p;p+=C;c.shw=p;p+=C;c.sha=p;p+=C;c.shg=p;p+=C;
    c.rvec=p;p+=C;c.kvec=p;p+=C;c.vvec=p;p+=C;c.sigw=p;p+=C;c.avec=p;p+=C;
    c.tk=p;p+=C;c.kk=p;p+=C;c.blend=p;p+=C;
    c.lw=p;p+=maxD;c.tw=p;p+=maxD;c.la=p;p+=maxD;c.lg=p;p+=maxD;c.sg=p;p+=maxD;c.lv=p;p+=maxD;

    for (int bch=0;bch<B;bch++) for (int t=0;t<T;t++){
        size_t base=((size_t)bch*T+t)*C;
        const float *xx=&x[base];
        const float *prev=(t==0)?(x_prev_init?&x_prev_init[(size_t)bch*C]:NULL)
                                :&x[(((size_t)bch*T)+(t-1))*C];
        float *pb=NULL; const float *prevp=prev;
        if(!prevp){pb=calloc(C,sizeof(float));prevp=pb;}
        c.gvec=&gv[base];
        tm_pre_forward(P,xx,prevp,&v_first[base],
            &R[base],&W[base],&K[base],&V[base],&A[base],&Bb[base],&c);
        free(pb);
    }
    rx_wkv7_forward_seq(R,W,K,V,A,Bb,NULL,y,NULL,B,T,H,N);

    for (int bch=0;bch<B;bch++) for (int t=0;t<T;t++){
        size_t base=((size_t)bch*T+t)*C;
        gn_forward(&y[base],P->ln_x_w,P->ln_x_b,ygn,mh,rh,H,N);
        for (int h=0;h<H;h++){int off=h*N;float s=0.0f;
            for(int j=0;j<N;j++) s+=R[base+off+j]*K[base+off+j]*P->r_k[off+j];
            for(int j=0;j<N;j++) ygn[off+j]+=s*V[base+off+j];}
        for (int i=0;i<C;i++) ygn[i]*=gv[base+i];      /* tg = y_bonus * g */
        mv_fwd(ygn,P->output_w,&out[base],C,C);
    }
    free(R);free(W);free(K);free(V);free(A);free(Bb);free(y);free(gv);
    free(cb);free(ygn);free(mh);free(rh);
    return 0;
}

/* ---- Backward ---- */

/* zero a grad buffer if present */
static inline void zpar(float *p, size_t n){ if (p) memset(p,0,n*sizeof(float)); }

int rx_timemix_backward(
    const rx_tm_params *P, const float *x, const float *x_prev_init,
    const float *v_first, const float *dy,
    rx_tm_grads *G, float *dx, float *d_xprev_init, float *d_v_first,
    int B, int T)
{
    int C = P->c, H = P->n_head, N = P->head_size;
    int Dw=P->Dw, Da=P->Da, Dg=P->Dg, Dv=(P->layer_id>0)?P->Dv:0;
    size_t BT=(size_t)B*T, tot=BT*C;

    /* ---- zero all outputs ----
     * d_v_first is an OUTPUT for layer>0 (grad w.r.t. the v_first *input*).  For
     * layer 0 (which PRODUCES v_first) it is instead an optional INPUT: the grad
     * arriving on the v_first output from downstream layers, injected into the
     * value path below.  So we only zero it when it is an output (layer>0);
     * callers that don't use v_first coupling pass NULL (legacy behaviour). */
    memset(dx,0,tot*sizeof(float));
    if (d_xprev_init) memset(d_xprev_init,0,(size_t)B*C*sizeof(float));
    if (d_v_first && P->layer_id>0) memset(d_v_first,0,tot*sizeof(float));
    zpar(G->x_r,C);zpar(G->x_k,C);zpar(G->x_v,C);zpar(G->x_w,C);zpar(G->x_a,C);zpar(G->x_g,C);
    zpar(G->receptance_w,(size_t)C*C);zpar(G->key_w,(size_t)C*C);
    zpar(G->value_w,(size_t)C*C);zpar(G->output_w,(size_t)C*C);
    zpar(G->w0,C);zpar(G->w1,(size_t)C*Dw);zpar(G->w2,(size_t)Dw*C);
    zpar(G->a0,C);zpar(G->a1,(size_t)C*Da);zpar(G->a2,(size_t)Da*C);
    zpar(G->g1,(size_t)C*Dg);zpar(G->g2,(size_t)Dg*C);
    if (P->layer_id>0){zpar(G->v0,C);zpar(G->v1,(size_t)C*Dv);zpar(G->v2,(size_t)Dv*C);}
    zpar(G->k_k,C);zpar(G->k_a,C);zpar(G->r_k,C);zpar(G->ln_x_w,C);zpar(G->ln_x_b,C);

    /* ---- per-(b,t) caches (full sequence) ---- */
    #define ALC(nm,sz) float *nm = malloc((size_t)(sz)*sizeof(float))
    ALC(R,tot);ALC(W,tot);ALC(K,tot);ALC(V,tot);ALC(A,tot);ALC(Bb,tot);ALC(y,tot);
    ALC(shr,tot);ALC(shk,tot);ALC(shv,tot);ALC(shw,tot);ALC(sha,tot);ALC(shg,tot);
    ALC(kvec,tot);ALC(vvec,tot);ALC(sigw,tot);ALC(avec,tot);ALC(gvec,tot);
    ALC(tk,tot);ALC(kk,tot);
    ALC(lw,BT*(size_t)Dw);ALC(tw,BT*(size_t)Dw);ALC(la,BT*(size_t)Da);
    ALC(lg,BT*(size_t)Dg);ALC(sg,BT*(size_t)Dg);
    ALC(lv,(Dv? BT*(size_t)Dv:1));ALC(blend,tot);
    ALC(ygn,tot);ALC(yb,tot);
    ALC(gn_mean,BT*(size_t)H);ALC(gn_rstd,BT*(size_t)H);ALC(bonus,BT*(size_t)H);
    ALC(dy_wkv,tot);
    ALC(dr_post,tot);ALC(dkmod_post,tot);ALC(dvf_post,tot);ALC(dg_post,tot);
    ALC(dR,tot);ALC(dW,tot);ALC(dK,tot);ALC(dV,tot);ALC(dA,tot);ALC(dBk,tot);
    float *vf = malloc(tot*sizeof(float)); /* local v_first (written at layer0) */
    #undef ALC
    /* NULL check (compact) */
    void *ptrs[] = {R,W,K,V,A,Bb,y,shr,shk,shv,shw,sha,shg,kvec,vvec,sigw,avec,
        gvec,tk,kk,lw,tw,la,lg,sg,lv,blend,ygn,yb,gn_mean,gn_rstd,bonus,dy_wkv,
        dr_post,dkmod_post,dvf_post,dg_post,dR,dW,dK,dV,dA,dBk,vf};
    for (unsigned i=0;i<sizeof(ptrs)/sizeof(ptrs[0]);i++) if(!ptrs[i]){
        for (unsigned j=0;j<sizeof(ptrs)/sizeof(ptrs[0]);j++) free(ptrs[j]);
        return -1;
    }
    if (P->layer_id>0) memcpy(vf, v_first, tot*sizeof(float));

    /* ================= PHASE 1: forward, fill caches ================= */
    tm_cache c;
    for (int bch=0;bch<B;bch++) for (int t=0;t<T;t++){
        size_t base=((size_t)bch*T+t)*C;
        const float *xx=&x[base];
        const float *prev=(t==0)?(x_prev_init?&x_prev_init[(size_t)bch*C]:NULL)
                                :&x[(((size_t)bch*T)+(t-1))*C];
        float *pb=NULL; const float *prevp=prev;
        if(!prevp){pb=calloc(C,sizeof(float));prevp=pb;}
        c.shr=&shr[base];c.shk=&shk[base];c.shv=&shv[base];
        c.shw=&shw[base];c.sha=&sha[base];c.shg=&shg[base];
        c.rvec=&R[base]; c.kvec=&kvec[base]; c.vvec=&vvec[base];
        c.sigw=&sigw[base]; c.avec=&avec[base]; c.gvec=&gvec[base];
        c.tk=&tk[base]; c.kk=&kk[base];
        c.lw=&lw[(size_t)(bch*T+t)*Dw]; c.tw=&tw[(size_t)(bch*T+t)*Dw];
        c.la=&la[(size_t)(bch*T+t)*Da];
        c.lg=&lg[(size_t)(bch*T+t)*Dg]; c.sg=&sg[(size_t)(bch*T+t)*Dg];
        c.lv=Dv?&lv[(size_t)(bch*T+t)*Dv]:lv; c.blend=&blend[base];
        tm_pre_forward(P, xx, prevp, &vf[base],
            &R[base],&W[base],&K[base],&V[base],&A[base],&Bb[base], &c);
        free(pb);
    }
    rx_wkv7_forward_seq(R,W,K,V,A,Bb,NULL,y,NULL,B,T,H,N);

    /* ============ PHASE 2: post-processing backward per token ========= */
    for (int bch=0;bch<B;bch++) for (int t=0;t<T;t++){
        size_t base=((size_t)bch*T+t)*C;
        float *mh=&gn_mean[(size_t)(bch*T+t)*H], *rh=&gn_rstd[(size_t)(bch*T+t)*H];
        float *bon=&bonus[(size_t)(bch*T+t)*H];
        /* forward: group-norm -> ygn, bonus -> yb, (g gate + output in bwd) */
        gn_forward(&y[base],P->ln_x_w,P->ln_x_b,&ygn[base],mh,rh,H,N);
        for (int i=0;i<C;i++) yb[base+i]=ygn[base+i];
        for (int h=0;h<H;h++){int off=h*N;float s=0.0f;
            for(int j=0;j<N;j++) s+=R[base+off+j]*K[base+off+j]*P->r_k[off+j];
            bon[h]=s;
            for(int j=0;j<N;j++) yb[base+off+j]+=s*V[base+off+j];
        }
        /* --- backward --- */
        const float *dout=&dy[base];
        float dtg[RX_TM_MAXC], dyb[RX_TM_MAXC], dygn[RX_TM_MAXC];
        /* output proj: out = tg @ Wo, with tg = yb * g */
        for (int i=0;i<C;i++){
            float tgi=yb[base+i]*gvec[base+i];
            float acc=0.0f; const float *Woi=&P->output_w[(size_t)i*C];
            float *dWoi=G->output_w?&G->output_w[(size_t)i*C]:NULL;
            for (int j=0;j<C;j++){ acc+=dout[j]*Woi[j]; if(dWoi) dWoi[j]+=tgi*dout[j]; }
            dtg[i]=acc;
        }
        /* split dtg -> dyb (=dygn seed) and dg */
        for (int i=0;i<C;i++){
            dyb[i]=dtg[i]*gvec[base+i];
            dg_post[base+i]=dtg[i]*yb[base+i];
            dygn[i]=dyb[i];             /* yb = ygn + bonus*V => dygn seeded by dyb */
            dr_post[base+i]=0.0f; dkmod_post[base+i]=0.0f; dvf_post[base+i]=0.0f;
        }
        /* bonus backward: yb[j] += bonus_h * V[j]; bonus_h = sum_j r*kmod*r_k */
        for (int h=0;h<H;h++){int off=h*N;
            float db=0.0f;
            for(int j=0;j<N;j++) db+=dyb[off+j]*V[base+off+j];
            for(int j=0;j<N;j++){
                float rj=R[base+off+j], kj=K[base+off+j], rkj=P->r_k[off+j];
                dvf_post[base+off+j]+=dyb[off+j]*bon[h];
                dr_post[base+off+j]  +=db*kj*rkj;
                dkmod_post[base+off+j]+=db*rj*rkj;
                if (G->r_k) G->r_k[off+j]+=db*rj*kj;
            }
        }
        /* group-norm backward -> dy_wkv (grad w.r.t. WKV7 output y) */
        for (int h=0;h<H;h++){int off=h*N;
            float mean=mh[h], rstd=rh[h];
            float mdn=0.0f, mdnn=0.0f, normv[RX_TM_MAXN];
            for(int j=0;j<N;j++){
                float nv=(y[base+off+j]-mean)*rstd; normv[j]=nv;
                if (G->ln_x_w) G->ln_x_w[off+j]+=dygn[off+j]*nv;
                if (G->ln_x_b) G->ln_x_b[off+j]+=dygn[off+j];
                float dnorm=dygn[off+j]*P->ln_x_w[off+j];
                mdn+=dnorm; mdnn+=dnorm*nv;
            }
            mdn/=N; mdnn/=N;
            for(int j=0;j<N;j++){
                float dnorm=dygn[off+j]*P->ln_x_w[off+j];
                dy_wkv[base+off+j]=rstd*(dnorm - mdn - normv[j]*mdnn);
            }
        }
    }

    /* ============ PHASE 3: WKV7 backward over the sequence =========== */
    rx_wkv7_backward(R,W,K,V,A,Bb,dy_wkv,NULL,
                     dR,dW,dK,dV,dA,dBk,NULL,B,T,H,N);

    /* ============ PHASE 4: preprocessing backward per token ========== */
    for (int bch=0;bch<B;bch++) for (int t=0;t<T;t++){
        size_t base=((size_t)bch*T+t)*C;
        const float *xx=&x[base];
        int tprev_valid = (t>0);
        const float *prev = tprev_valid ? &x[(((size_t)bch*T)+(t-1))*C]
                          : (x_prev_init?&x_prev_init[(size_t)bch*C]:NULL);
        float *pb=NULL; const float *prevp=prev;
        if(!prevp){pb=calloc(C,sizeof(float));prevp=pb;}
        /* destinations for dprev */
        float *dprev_dst = tprev_valid ? &dx[(((size_t)bch*T)+(t-1))*C]
                        : (d_xprev_init?&d_xprev_init[(size_t)bch*C]:NULL);

        size_t td=(size_t)(bch*T+t);
        float dr_t[RX_TM_MAXC],dkmod_t[RX_TM_MAXC],dvf_t[RX_TM_MAXC];
        float dkvec[RX_TM_MAXC],dvvec[RX_TM_MAXC],dav[RX_TM_MAXC],dkk[RX_TM_MAXC],dtk[RX_TM_MAXC];
        float dsh_r[RX_TM_MAXC],dsh_k[RX_TM_MAXC],dsh_v[RX_TM_MAXC],dsh_w[RX_TM_MAXC],dsh_a[RX_TM_MAXC],dsh_g[RX_TM_MAXC];
        float tmpc[RX_TM_MAXC];
        for (int i=0;i<C;i++){
            dr_t[i]=dR[base+i]+dr_post[base+i];
            dkmod_t[i]=dK[base+i]+dkmod_post[base+i];
            dvf_t[i]=dV[base+i]+dvf_post[base+i];
            dkvec[i]=0.0f; dvvec[i]=0.0f; dav[i]=0.0f;
            dsh_r[i]=0.0f;dsh_k[i]=0.0f;dsh_v[i]=0.0f;
            dsh_w[i]=0.0f;dsh_a[i]=0.0f;dsh_g[i]=0.0f;
        }
        /* recombine nkk/kka -> dkk, dav */
        for (int i=0;i<C;i++){
            float dnkk=dA[base+i], dkka=dBk[base+i];
            dkk[i]=-dnkk + avec[base+i]*dkka;
            dav[i]+=kk[base+i]*dkka;
        }
        /* k_mod backward: K = kvec*(1+(a-1)*k_a) */
        for (int i=0;i<C;i++){
            float av=avec[base+i], ka=P->k_a[i], kv=kvec[base+i];
            float factor=1.0f+(av-1.0f)*ka;
            dkvec[i]+=dkmod_t[i]*factor;
            dav[i]+=dkmod_t[i]*kv*ka;
            if (G->k_a) G->k_a[i]+=dkmod_t[i]*kv*(av-1.0f);
        }
        /* kk L2-norm backward per head -> dtk.  The radial-projection term
         * (inv^2*gdot*u/norm) nearly cancels inv*dkk in fp32; accumulate the
         * reductions and combine in double to keep k_k / kk-path grads clean
         * (the L2-normalize analogue of the a/b noise floor, ref §6). */
        for (int h=0;h<H;h++){int off=h*N;
            double ss=0.0; for(int j=0;j<N;j++){double u=tk[base+off+j]; ss+=u*u;}
            double norm=sqrt(ss), inv=1.0/(norm+1e-12);
            double gdot=0.0; for(int j=0;j<N;j++) gdot+=(double)dkk[off+j]*(double)tk[base+off+j];
            for(int j=0;j<N;j++){
                double u=tk[base+off+j];
                double d = inv*(double)dkk[off+j];
                if (norm>0.0) d -= inv*inv*gdot*u/norm;
                dtk[off+j]=(float)d;
            }
        }
        /* tk = kvec*k_k */
        for (int i=0;i<C;i++){
            dkvec[i]+=dtk[i]*P->k_k[i];
            if (G->k_k) G->k_k[i]+=dtk[i]*kvec[base+i];
        }
        /* a_vec = sigmoid(a0+oa) */
        for (int i=0;i<C;i++){
            float av=avec[base+i];
            tmpc[i]=dav[i]*av*(1.0f-av);           /* d_oa (=dsa_) */
            if (G->a0) G->a0[i]+=tmpc[i];
        }
        { float d_la[RX_TM_MAXC];
          for(int d=0;d<Da;d++) d_la[d]=0.0f;
          mv_bwd(&la[td*Da], P->a2, tmpc, d_la, G->a2, Da, C);
          mv_bwd(&sha[base], P->a1, d_la, dsh_a, G->a1, C, Da);
        }
        /* w_vec = exp(-coeff*sigmoid(w0+ow)) */
        for (int i=0;i<C;i++){
            float wv=W[base+i], sig=sigw[base+i];
            float dsig=dW[base+i]*wv*(-RX_WDECAY_COEFF);
            float dsw=dsig*sig*(1.0f-sig);
            tmpc[i]=dsw; /* d_ow */
            if (G->w0) G->w0[i]+=dsw;
        }
        { float d_tw[RX_TM_MAXC];
          for(int d=0;d<Dw;d++) d_tw[d]=0.0f;
          mv_bwd(&tw[td*Dw], P->w2, tmpc, d_tw, G->w2, Dw, C);
          float d_lw[RX_TM_MAXC];
          for(int d=0;d<Dw;d++){ float tv=tw[td*Dw+d]; d_lw[d]=d_tw[d]*(1.0f-tv*tv); }
          mv_bwd(&shw[base], P->w1, d_lw, dsh_w, G->w1, C, Dw);
        }
        /* g_vec = sg@g2 ; sg=sigmoid(lg) */
        { float d_sg[RX_TM_MAXC];
          for(int d=0;d<Dg;d++) d_sg[d]=0.0f;
          mv_bwd(&sg[td*Dg], P->g2, &dg_post[base], d_sg, G->g2, Dg, C);
          float d_lg[RX_TM_MAXC];
          for(int d=0;d<Dg;d++){ float sv=sg[td*Dg+d]; d_lg[d]=d_sg[d]*sv*(1.0f-sv); }
          mv_bwd(&shg[base], P->g1, d_lg, dsh_g, G->g1, C, Dg);
        }
        /* value residual */
        if (P->layer_id==0){
            /* vf == v at layer 0; add the block's own dv plus any downstream
             * grad on the produced v_first (d_v_first as INPUT, see above). */
            for (int i=0;i<C;i++){
                dvvec[i]+=dvf_t[i];
                if (d_v_first) dvvec[i]+=d_v_first[base+i];
            }
        } else {
            for (int i=0;i<C;i++){
                float bl=blend[base+i];
                dvvec[i]+=dvf_t[i]*(1.0f-bl);
                if (d_v_first) d_v_first[base+i]+=dvf_t[i]*bl;
                float dbl=dvf_t[i]*(vf[base+i]-vvec[base+i]);
                float dsv=dbl*bl*(1.0f-bl);
                tmpc[i]=dsv; /* d_ov */
                if (G->v0) G->v0[i]+=dsv;
            }
            float d_lv[RX_TM_MAXC];
            for(int d=0;d<Dv;d++) d_lv[d]=0.0f;
            mv_bwd(&lv[td*Dv], P->v2, tmpc, d_lv, G->v2, Dv, C);
            mv_bwd(&shv[base], P->v1, d_lv, dsh_v, G->v1, C, Dv);
        }
        /* projections: r,k,v */
        mv_bwd(&shr[base], P->receptance_w, dr_t,  dsh_r, G->receptance_w, C, C);
        mv_bwd(&shk[base], P->key_w,        dkvec, dsh_k, G->key_w,        C, C);
        mv_bwd(&shv[base], P->value_w,      dvvec, dsh_v, G->value_w,      C, C);

        /* token-shift backward for all six lerps */
        #define TSHIFT(dsh, lerp, glerp) do { \
            for (int i=0;i<C;i++){ \
                float dd=dsh[i]; \
                if (glerp) glerp[i]+=dd*(prevp[i]-xx[i]); \
                dx[base+i]+=dd*(1.0f-lerp[i]); \
                if (dprev_dst) dprev_dst[i]+=dd*lerp[i]; \
            } } while(0)
        TSHIFT(dsh_r, P->x_r, G->x_r);
        TSHIFT(dsh_k, P->x_k, G->x_k);
        TSHIFT(dsh_v, P->x_v, G->x_v);
        TSHIFT(dsh_w, P->x_w, G->x_w);
        TSHIFT(dsh_a, P->x_a, G->x_a);
        TSHIFT(dsh_g, P->x_g, G->x_g);
        #undef TSHIFT
        free(pb);
    }

    for (unsigned j=0;j<sizeof(ptrs)/sizeof(ptrs[0]);j++) free(ptrs[j]);
    return 0;
}

/* ================================================================
 * RWKV-7 CHANNEL-MIX (ReLU^2 FFN) forward/backward.
 *
 * Mirrors rwkvx.c `rwkv7_channel_mix_fast`:
 *   shifted = xx + (prev - xx) * ffn_x_k          (token shift / lerp)
 *   fk      = shifted @ ffn_key_w                 ([c] -> [ffn_dim])
 *   k       = relu(fk)^2                          (squared ReLU)
 *   out     = k @ ffn_value_w                     ([ffn_dim] -> [c])
 * prev at token t is the previous token's channel-mix input (ln2 output);
 * prev at t=0 is x_prev_init (or zero).
 * ================================================================ */
int rx_channelmix_forward_seq(
    const rx_cm_params *P, const float *x, const float *x_prev_init,
    float *out, int B, int T)
{
    int C = P->c, F = P->ffn_dim;
    float *sh = malloc((size_t)C*sizeof(float));
    float *fk = malloc((size_t)F*sizeof(float));
    if (!sh || !fk) { free(sh); free(fk); return -1; }
    for (int bch=0;bch<B;bch++) for (int t=0;t<T;t++){
        size_t base=((size_t)bch*T+t)*C;
        const float *xx=&x[base];
        const float *prev=(t==0)?(x_prev_init?&x_prev_init[(size_t)bch*C]:NULL)
                                :&x[(((size_t)bch*T)+(t-1))*C];
        for (int i=0;i<C;i++){ float pv=prev?prev[i]:0.0f;
            sh[i]=xx[i]+(pv-xx[i])*P->ffn_x_k[i]; }
        mv_fwd(sh, P->ffn_key_w, fk, C, F);
        for (int m=0;m<F;m++){ float v=fk[m]>0.0f?fk[m]:0.0f; fk[m]=v*v; }
        mv_fwd(fk, P->ffn_value_w, &out[base], F, C);
    }
    free(sh); free(fk);
    return 0;
}

int rx_channelmix_backward(
    const rx_cm_params *P, const float *x, const float *x_prev_init,
    const float *dy, rx_cm_grads *G, float *dx, float *d_xprev_init,
    int B, int T)
{
    int C = P->c, F = P->ffn_dim;
    size_t tot=(size_t)B*T*C;
    memset(dx,0,tot*sizeof(float));
    if (d_xprev_init) memset(d_xprev_init,0,(size_t)B*C*sizeof(float));
    zpar(G->ffn_x_k,C);
    zpar(G->ffn_key_w,(size_t)C*F);
    zpar(G->ffn_value_w,(size_t)F*C);

    float *sh=malloc((size_t)C*sizeof(float));
    float *fk=malloc((size_t)F*sizeof(float));   /* pre-activation */
    float *ka=malloc((size_t)F*sizeof(float));   /* relu(fk)^2 */
    float *dfk=malloc((size_t)F*sizeof(float));
    float *dsh=malloc((size_t)C*sizeof(float));
    if (!sh||!fk||!ka||!dfk||!dsh){ free(sh);free(fk);free(ka);free(dfk);free(dsh);return -1; }

    for (int bch=0;bch<B;bch++) for (int t=0;t<T;t++){
        size_t base=((size_t)bch*T+t)*C;
        const float *xx=&x[base];
        int tprev_valid=(t>0);
        const float *prev=tprev_valid?&x[(((size_t)bch*T)+(t-1))*C]
                          :(x_prev_init?&x_prev_init[(size_t)bch*C]:NULL);
        float *dprev_dst=tprev_valid?&dx[(((size_t)bch*T)+(t-1))*C]
                        :(d_xprev_init?&d_xprev_init[(size_t)bch*C]:NULL);
        /* recompute forward for this token */
        for (int i=0;i<C;i++){ float pv=prev?prev[i]:0.0f;
            sh[i]=xx[i]+(pv-xx[i])*P->ffn_x_k[i]; }
        mv_fwd(sh, P->ffn_key_w, fk, C, F);
        for (int m=0;m<F;m++){ float v=fk[m]>0.0f?fk[m]:0.0f; ka[m]=v*v; }

        /* value proj backward: out = ka @ Wv */
        for (int m=0;m<F;m++) dfk[m]=0.0f;
        mv_bwd(ka, P->ffn_value_w, &dy[base], dfk, G->ffn_value_w, F, C);
        /* squared-relu backward: k = relu(fk)^2 => dfk = dk * 2*relu(fk) */
        for (int m=0;m<F;m++){ float r=fk[m]>0.0f?fk[m]:0.0f; dfk[m]*=2.0f*r; }
        /* key proj backward: fk = sh @ Wk */
        for (int i=0;i<C;i++) dsh[i]=0.0f;
        mv_bwd(sh, P->ffn_key_w, dfk, dsh, G->ffn_key_w, C, F);
        /* token-shift backward: sh = xx + (prev-xx)*x_k */
        for (int i=0;i<C;i++){
            float dd=dsh[i], pv=prev?prev[i]:0.0f;
            if (G->ffn_x_k) G->ffn_x_k[i]+=dd*(pv-xx[i]);
            dx[base+i]+=dd*(1.0f-P->ffn_x_k[i]);
            if (dprev_dst) dprev_dst[i]+=dd*P->ffn_x_k[i];
        }
    }
    free(sh);free(fk);free(ka);free(dfk);free(dsh);
    return 0;
}

/* ================================================================
 * LayerNorm forward/backward (eps=1e-5, over full C; foundry fast_layer_norm).
 * Mean/variance reductions accumulated in DOUBLE to kill the fp32 cancellation
 * (same discipline as the L2-norm radial term; see §6).
 * ================================================================ */
#define RX_LN_EPS 1e-5

int rx_layernorm_forward_seq(
    const float *x, const float *w, const float *b, float *out,
    int B, int T, int C)
{
    for (int bch=0;bch<B;bch++) for (int t=0;t<T;t++){
        size_t base=((size_t)bch*T+t)*C;
        const float *xx=&x[base];
        double mean=0.0; for (int i=0;i<C;i++) mean+=xx[i]; mean/=C;
        double var=0.0; for (int i=0;i<C;i++){ double d=xx[i]-mean; var+=d*d; } var/=C;
        double rstd=1.0/sqrt(var+RX_LN_EPS);
        for (int i=0;i<C;i++)
            out[base+i]=(float)(((double)xx[i]-mean)*rstd)*w[i]+b[i];
    }
    return 0;
}

int rx_layernorm_backward(
    const float *x, const float *w, const float *b, const float *dy,
    float *dx, float *dw, float *db, int B, int T, int C)
{
    (void)b;
    if (dw) memset(dw,0,(size_t)C*sizeof(float));
    if (db) memset(db,0,(size_t)C*sizeof(float));
    float *norm=malloc((size_t)C*sizeof(float));
    if (!norm) return -1;
    for (int bch=0;bch<B;bch++) for (int t=0;t<T;t++){
        size_t base=((size_t)bch*T+t)*C;
        const float *xx=&x[base];
        double mean=0.0; for (int i=0;i<C;i++) mean+=xx[i]; mean/=C;
        double var=0.0; for (int i=0;i<C;i++){ double d=xx[i]-mean; var+=d*d; } var/=C;
        double rstd=1.0/sqrt(var+RX_LN_EPS);
        /* reductions of dnorm and dnorm*norm in double */
        double mdn=0.0, mdnn=0.0;
        for (int i=0;i<C;i++){
            double nv=((double)xx[i]-mean)*rstd; norm[i]=(float)nv;
            double dnorm=(double)dy[base+i]*(double)w[i];
            mdn+=dnorm; mdnn+=dnorm*nv;
            if (dw) dw[i]+=dy[base+i]*(float)nv;
            if (db) db[i]+=dy[base+i];
        }
        mdn/=C; mdnn/=C;
        for (int i=0;i<C;i++){
            double dnorm=(double)dy[base+i]*(double)w[i];
            dx[base+i]=(float)(rstd*(dnorm - mdn - (double)norm[i]*mdnn));
        }
    }
    free(norm);
    return 0;
}

/* ================================================================
 * RWKV-7 BLOCK forward/backward.
 *   h1 = ln1(x);   tm = timemix(h1);   x1 = x + tm
 *   h2 = ln2(x1);  cm = channelmix(h2); out = x1 + cm
 * ================================================================ */
int rx_block_forward_seq(
    const rx_block_params *P, const float *x,
    const float *tm_xprev_init, const float *cm_xprev_init,
    float *v_first, float *out, int B, int T)
{
    int C=P->tm.c; size_t tot=(size_t)B*T*C;
    float *h1=malloc(tot*sizeof(float));
    float *tm=malloc(tot*sizeof(float));
    float *x1=malloc(tot*sizeof(float));
    float *h2=malloc(tot*sizeof(float));
    float *cm=malloc(tot*sizeof(float));
    if (!h1||!tm||!x1||!h2||!cm){ free(h1);free(tm);free(x1);free(h2);free(cm);return -1; }
    int rc=0;
    rc|=rx_layernorm_forward_seq(x,P->ln1_w,P->ln1_b,h1,B,T,C);
    rc|=rx_timemix_forward_seq(&P->tm,h1,tm_xprev_init,v_first,tm,B,T);
    for (size_t i=0;i<tot;i++) x1[i]=x[i]+tm[i];
    rc|=rx_layernorm_forward_seq(x1,P->ln2_w,P->ln2_b,h2,B,T,C);
    rc|=rx_channelmix_forward_seq(&P->cm,h2,cm_xprev_init,cm,B,T);
    for (size_t i=0;i<tot;i++) out[i]=x1[i]+cm[i];
    free(h1);free(tm);free(x1);free(h2);free(cm);
    return rc?-1:0;
}

int rx_block_backward(
    const rx_block_params *P, const float *x,
    const float *tm_xprev_init, const float *cm_xprev_init,
    const float *v_first, const float *dy,
    rx_block_grads *G, float *dx,
    float *d_tm_xprev_init, float *d_cm_xprev_init, float *d_v_first,
    int B, int T)
{
    int C=P->tm.c; size_t tot=(size_t)B*T*C;
    /* forward temporaries needed by the sub-backwards */
    float *h1=malloc(tot*sizeof(float));
    float *tm=malloc(tot*sizeof(float));
    float *x1=malloc(tot*sizeof(float));
    float *h2=malloc(tot*sizeof(float));
    float *vf=malloc(tot*sizeof(float));   /* v_first (written at layer 0) */
    /* backward temporaries */
    float *dx1=malloc(tot*sizeof(float));
    float *dh2=malloc(tot*sizeof(float));
    float *dh1=malloc(tot*sizeof(float));
    float *dxa=malloc(tot*sizeof(float));  /* dx contribution from an LN */
    void *ps[]={h1,tm,x1,h2,vf,dx1,dh2,dh1,dxa};
    for (unsigned i=0;i<sizeof(ps)/sizeof(ps[0]);i++) if(!ps[i]){
        for (unsigned j=0;j<sizeof(ps)/sizeof(ps[0]);j++) free(ps[j]); return -1; }
    if (P->tm.layer_id>0 && v_first) memcpy(vf,v_first,tot*sizeof(float));

    /* ---- forward ---- */
    rx_layernorm_forward_seq(x,P->ln1_w,P->ln1_b,h1,B,T,C);
    rx_timemix_forward_seq(&P->tm,h1,tm_xprev_init,vf,tm,B,T);
    for (size_t i=0;i<tot;i++) x1[i]=x[i]+tm[i];
    rx_layernorm_forward_seq(x1,P->ln2_w,P->ln2_b,h2,B,T,C);

    /* ---- backward (reverse) ---- */
    /* out = x1 + cm  =>  d_cm = dy,  dx1 seeded with dy (skip) */
    /* channel-mix: d_cm -> dh2 + cm param grads */
    rx_channelmix_backward(&P->cm,h2,cm_xprev_init,dy,&G->cm,dh2,d_cm_xprev_init,B,T);
    /* ln2: dh2 -> dxa (=d(x1) through the norm) + ln2 grads */
    rx_layernorm_backward(x1,P->ln2_w,P->ln2_b,dh2,dxa,G->ln2_w,G->ln2_b,B,T,C);
    for (size_t i=0;i<tot;i++) dx1[i]=dy[i]+dxa[i];

    /* x1 = x + tm  =>  d_tm = dx1,  dx seeded with dx1 (skip) */
    rx_timemix_backward(&P->tm,h1,tm_xprev_init,vf,dx1,&G->tm,dh1,
                        d_tm_xprev_init,d_v_first,B,T);
    /* ln1: dh1 -> dxa (=d(x) through the norm) + ln1 grads */
    rx_layernorm_backward(x,P->ln1_w,P->ln1_b,dh1,dxa,G->ln1_w,G->ln1_b,B,T,C);
    for (size_t i=0;i<tot;i++) dx[i]=dx1[i]+dxa[i];

    for (unsigned j=0;j<sizeof(ps)/sizeof(ps[0]);j++) free(ps[j]);
    return 0;
}

/* ================================================================
 * SPARSE (MoBA-style) ATTENTION BLOCK forward/backward — BASE weights.
 * Mirrors rwkvx.c `attn_block_forward_fast` (the differentiable core):
 *   h1 = ln1(x); q,k,v = h1@{Wq,Wk,Wv};
 *   att = MoBA-sparse causal softmax-attn(q,k,v); att_out = att@Wo;
 *   x1 = x + att_out; h2 = ln2(x1); cm = channelmix(h2); out = x1 + cm
 *
 * MoBA top-k chunk selection is HARD: the selected key set is treated as fixed;
 * unselected keys carry attention weight exactly 0, so they vanish from every
 * backward reduction with no special-casing.  See rx_pretrain.h.
 * ================================================================ */

/* Attention core: fills h1, q, k, v [B*T*C], softmax weights p
 * [B*n_head*T*T] and attention output y [B*T*C].  Selection + softmax are
 * evaluated in double.  p[s]=0 for keys that are non-causal (s>tq) or in an
 * unselected chunk. */
static void attn_score_av(const rx_attn_params *P, const float *q, const float *k,
                          const float *v, int B, int T, float *p, float *y);
static void attn_score_av_bwd(const rx_attn_params *P, const float *q, const float *k,
                          const float *v, const float *p, const float *dyatt,
                          int B, int T, float *dq, float *dk, float *dv);
static void rx_attn_core(const rx_attn_params *P, const float *x, int B, int T,
                         float *h1, float *q, float *k, float *v,
                         float *p, float *y)
{
    int C=P->c;
    rx_layernorm_forward_seq(x, P->ln1_w, P->ln1_b, h1, B, T, C);
    for (int b=0;b<B;b++) for (int t=0;t<T;t++){
        size_t base=((size_t)b*T+t)*C;
        mv_fwd(&h1[base], P->receptance_w, &q[base], C, C);
        mv_fwd(&h1[base], P->key_w,        &k[base], C, C);
        mv_fwd(&h1[base], P->value_w,      &v[base], C, C);
    }
    attn_score_av(P,q,k,v,B,T,p,y);
}

/* Score/softmax/AV attention core: given projected q,k,v [B*T*C], fill the
 * MoBA softmax weights p [B*n_head*T*T] and attention output y [B*T*C].  Chunk
 * selection + softmax reductions are done in double.  Extracted verbatim from
 * the old rx_attn_core so the GPU MoBA kernel (cuda_kernels/moba_attn.cu) is
 * parity-checked against exactly this reduction. */
static void attn_score_av(const rx_attn_params *P, const float *q, const float *k,
                          const float *v, int B, int T, float *p, float *y)
{
    int C=P->c, H=P->n_head, N=P->head_size;
    int cs = P->chunk_size<1 ? 1 : P->chunk_size;
    float scale = 1.0f/sqrtf((float)N);
    for (int b=0;b<B;b++) for (int h=0;h<H;h++){
        int off=h*N;
        for (int tq=0;tq<T;tq++){
            size_t qb=((size_t)b*T+tq)*C + off;
            float *pv = &p[(((size_t)(b*H+h))*T + tq)*T];
            for (int s=0;s<T;s++) pv[s]=0.0f;

            int local_chunk = tq/cs;
            int ncand = local_chunk;               /* earlier complete chunks */
            unsigned char csel[T>0?T:1];
            int keepall = (P->moba_top_k<=0 || P->moba_top_k>=ncand);
            if (keepall){
                for (int cc=0;cc<ncand;cc++) csel[cc]=1;
            } else {
                double gate[T>0?T:1];
                for (int cc=0;cc<ncand;cc++){
                    double g=0.0;
                    for (int j=0;j<N;j++){
                        double mkj=0.0;
                        for (int s=cc*cs;s<(cc+1)*cs;s++)
                            mkj += (double)k[((size_t)b*T+s)*C+off+j];
                        mkj/=(double)cs;
                        g += (double)q[qb+j]*mkj;
                    }
                    gate[cc]=g; csel[cc]=0;
                }
                for (int it=0; it<P->moba_top_k; it++){
                    int arg=-1; double best=-1e300;
                    for (int cc=0;cc<ncand;cc++)
                        if (!csel[cc] && gate[cc]>best){ best=gate[cc]; arg=cc; }
                    if (arg<0) break; csel[arg]=1;
                }
            }

            /* scores over selected causal keys (double softmax) */
            double sc[T>0?T:1];
            double maxs=-1e300;
            for (int s=0;s<=tq;s++){
                int cc=s/cs;
                int inc = (cc==local_chunk) || (cc<ncand && csel[cc]);
                if (!inc){ sc[s]=-1e300; continue; }
                double d=0.0;
                for (int j=0;j<N;j++)
                    d += (double)q[qb+j]*(double)k[((size_t)b*T+s)*C+off+j];
                d*=scale; sc[s]=d; if (d>maxs) maxs=d;
            }
            double sum=0.0;
            for (int s=0;s<=tq;s++){
                if (sc[s]<=-1e299) continue;
                double e=exp(sc[s]-maxs); sc[s]=e; sum+=e;
            }
            for (int s=0;s<=tq;s++){
                if (sc[s]<=-1e299) continue;
                pv[s]=(float)(sc[s]/sum);
            }
            for (int j=0;j<N;j++){
                double acc=0.0;
                for (int s=0;s<=tq;s++)
                    if (pv[s]!=0.0f)
                        acc += (double)pv[s]*(double)v[((size_t)b*T+s)*C+off+j];
                y[qb+j]=(float)acc;
            }
        }
    }
}

int rx_attn_forward_seq(
    const rx_attn_params *P, const float *x, const float *cm_xprev_init,
    float *out, int B, int T)
{
    int C=P->c, H=P->n_head; size_t tot=(size_t)B*T*C;
    float *h1=malloc(tot*sizeof(float));
    float *q =malloc(tot*sizeof(float));
    float *k =malloc(tot*sizeof(float));
    float *v =malloc(tot*sizeof(float));
    float *y =malloc(tot*sizeof(float));
    float *x1=malloc(tot*sizeof(float));
    float *h2=malloc(tot*sizeof(float));
    float *cm=malloc(tot*sizeof(float));
    float *p =malloc((size_t)B*H*T*T*sizeof(float));
    void *ps[]={h1,q,k,v,y,x1,h2,cm,p};
    for (unsigned i=0;i<sizeof(ps)/sizeof(ps[0]);i++) if(!ps[i]){
        for (unsigned j=0;j<sizeof(ps)/sizeof(ps[0]);j++) free(ps[j]); return -1; }

    rx_attn_core(P,x,B,T,h1,q,k,v,p,y);
    for (int b=0;b<B;b++) for (int t=0;t<T;t++){
        size_t base=((size_t)b*T+t)*C;
        mv_fwd(&y[base], P->output_w, &x1[base], C, C);  /* att_out into x1 */
        for (int i=0;i<C;i++) x1[base+i]+=x[base+i];      /* residual */
    }
    rx_layernorm_forward_seq(x1,P->ln2_w,P->ln2_b,h2,B,T,C);
    rx_channelmix_forward_seq(&P->cm,h2,cm_xprev_init,cm,B,T);
    for (size_t i=0;i<tot;i++) out[i]=x1[i]+cm[i];

    for (unsigned j=0;j<sizeof(ps)/sizeof(ps[0]);j++) free(ps[j]);
    return 0;
}

int rx_attn_backward_full(
    const rx_attn_params *P, const float *x, const float *cm_xprev_init,
    const float *dy, rx_attn_grads *G, float *dx, float *d_cm_xprev_init,
    int B, int T)
{
    int C=P->c, H=P->n_head; size_t tot=(size_t)B*T*C;

    /* zero the base-projection grads accumulated via mv_bwd (ln/cm grads are
     * zeroed by their own sub-backwards) */
    zpar(G->receptance_w,(size_t)C*C); zpar(G->key_w,(size_t)C*C);
    zpar(G->value_w,(size_t)C*C);      zpar(G->output_w,(size_t)C*C);

    /* forward temporaries */
    float *h1=malloc(tot*sizeof(float));
    float *q =malloc(tot*sizeof(float));
    float *k =malloc(tot*sizeof(float));
    float *v =malloc(tot*sizeof(float));
    float *y =malloc(tot*sizeof(float));
    float *x1=malloc(tot*sizeof(float));
    float *h2=malloc(tot*sizeof(float));
    float *p =malloc((size_t)B*H*T*T*sizeof(float));
    /* backward temporaries */
    float *dh2=malloc(tot*sizeof(float));
    float *dxa=malloc(tot*sizeof(float));
    float *dx1=malloc(tot*sizeof(float));
    float *dyatt=malloc(tot*sizeof(float));   /* d(y) */
    float *dq=malloc(tot*sizeof(float));
    float *dk=malloc(tot*sizeof(float));
    float *dv=malloc(tot*sizeof(float));
    float *dh1=malloc(tot*sizeof(float));
    void *ps[]={h1,q,k,v,y,x1,h2,p,dh2,dxa,dx1,dyatt,dq,dk,dv,dh1};
    for (unsigned i=0;i<sizeof(ps)/sizeof(ps[0]);i++) if(!ps[i]){
        for (unsigned j=0;j<sizeof(ps)/sizeof(ps[0]);j++) free(ps[j]); return -1; }

    /* ---- recompute forward (selection fixed at base weights) ---- */
    rx_attn_core(P,x,B,T,h1,q,k,v,p,y);
    for (int b=0;b<B;b++) for (int t=0;t<T;t++){
        size_t base=((size_t)b*T+t)*C;
        mv_fwd(&y[base], P->output_w, &x1[base], C, C);
        for (int i=0;i<C;i++) x1[base+i]+=x[base+i];
    }
    rx_layernorm_forward_seq(x1,P->ln2_w,P->ln2_b,h2,B,T,C);  /* h2 = ln2(x1) */

    /* ---- backward (reverse) ---- */
    /* out = x1 + cm  =>  d_cm = dy,  dx1 seeded with dy (skip) */
    /* channel-mix: d_cm(=dy) -> dh2 + cm grads */
    rx_channelmix_backward(&P->cm,h2,cm_xprev_init,dy,&G->cm,dh2,d_cm_xprev_init,B,T);
    /* ln2: dh2 -> dxa (=d(x1) through the norm) + ln2 grads */
    rx_layernorm_backward(x1,P->ln2_w,P->ln2_b,dh2,dxa,G->ln2_w,G->ln2_b,B,T,C);
    for (size_t i=0;i<tot;i++) dx1[i]=dy[i]+dxa[i];

    /* x1 = x + att_out  =>  d_att_out = dx1,  dx seeded with dx1 */
    /* output proj: att_out = y @ Wo  =>  d(y)=Wo·d_att_out, dWo += y⊗d_att_out */
    memset(dyatt,0,tot*sizeof(float));
    for (int b=0;b<B;b++) for (int t=0;t<T;t++){
        size_t base=((size_t)b*T+t)*C;
        mv_bwd(&y[base], P->output_w, &dx1[base], &dyatt[base],
               G->output_w, C, C);
    }

    /* attention backward: d(y) -> dq, dk, dv (selection fixed via p) */
    attn_score_av_bwd(P,q,k,v,p,dyatt,B,T,dq,dk,dv);

    /* q,k,v = h1 @ {Wq,Wk,Wv}  =>  dh1 accumulate + proj grads */
    memset(dh1,0,tot*sizeof(float));
    for (int b=0;b<B;b++) for (int t=0;t<T;t++){
        size_t base=((size_t)b*T+t)*C;
        mv_bwd(&h1[base], P->receptance_w, &dq[base], &dh1[base], G->receptance_w, C, C);
        mv_bwd(&h1[base], P->key_w,        &dk[base], &dh1[base], G->key_w,        C, C);
        mv_bwd(&h1[base], P->value_w,      &dv[base], &dh1[base], G->value_w,      C, C);
    }

    /* ln1: dh1 -> dxa (=d(x) through the norm) + ln1 grads; dx = dx1 + dxa */
    rx_layernorm_backward(x,P->ln1_w,P->ln1_b,dh1,dxa,G->ln1_w,G->ln1_b,B,T,C);
    for (size_t i=0;i<tot;i++) dx[i]=dx1[i]+dxa[i];

    for (unsigned j=0;j<sizeof(ps)/sizeof(ps[0]);j++) free(ps[j]);
    return 0;
}

/* Score/softmax/AV attention backward: given projected q,k,v, the fixed MoBA
 * softmax weights p, and d(y) [B*T*C], accumulate dq,dk,dv [B*T*C] (all zeroed
 * at entry).  The score-gradient reduction wsum is done in double.  Extracted
 * verbatim from rx_attn_backward_full so the GPU kernel is parity-checked
 * against exactly this reduction. */
static void attn_score_av_bwd(const rx_attn_params *P, const float *q, const float *k,
                          const float *v, const float *p, const float *dyatt,
                          int B, int T, float *dq, float *dk, float *dv)
{
    int C=P->c, H=P->n_head, N=P->head_size; size_t tot=(size_t)B*T*C;
    float scale = 1.0f/sqrtf((float)N);
    memset(dq,0,tot*sizeof(float));
    memset(dk,0,tot*sizeof(float));
    memset(dv,0,tot*sizeof(float));
    for (int b=0;b<B;b++) for (int h=0;h<H;h++){
        int off=h*N;
        for (int tq=0;tq<T;tq++){
            size_t qb=((size_t)b*T+tq)*C+off;
            const float *pv=&p[(((size_t)(b*H+h))*T+tq)*T];
            /* dp[s] = <d(y_tq), v_s> ; wsum = sum_s p[s]*dp[s]  (double) */
            double dp[T>0?T:1];
            double wsum=0.0;
            for (int s=0;s<=tq;s++){
                if (pv[s]==0.0f){ dp[s]=0.0; continue; }
                double d=0.0;
                size_t sb=((size_t)b*T+s)*C+off;
                for (int j=0;j<N;j++) d += (double)dyatt[qb+j]*(double)v[sb+j];
                dp[s]=d; wsum += (double)pv[s]*d;
            }
            for (int s=0;s<=tq;s++){
                if (pv[s]==0.0f) continue;
                size_t sb=((size_t)b*T+s)*C+off;
                double dscore = (double)pv[s]*(dp[s]-wsum);  /* softmax bwd */
                float dsc = scale*(float)dscore;
                for (int j=0;j<N;j++){
                    dq[qb+j]   += dsc*k[sb+j];
                    dk[sb+j]   += dsc*q[qb+j];
                    dv[sb+j]   += pv[s]*dyatt[qb+j];
                }
            }
        }
    }
}

/* ---- Exported test hooks: the ISOLATED attention core (scores->softmax->AV),
 * the "last hard kernel" mirrored on GPU in cuda_kernels/moba_attn.cu.  These
 * call exactly the same static routines rx_attn_forward_seq/backward_full use,
 * so the whole-block gradient test (test_attn_bwd) validates them. ---- */
int rx_attn_score_av_forward(const rx_attn_params *P, const float *q, const float *k,
    const float *v, int B, int T, float *p, float *y)
{
    attn_score_av(P,q,k,v,B,T,p,y);
    return 0;
}

int rx_attn_score_av_backward(const rx_attn_params *P, const float *q, const float *k,
    const float *v, const float *dy, int B, int T,
    float *dq, float *dk, float *dv)
{
    int H=P->n_head; size_t tot=(size_t)B*T*P->c;
    float *p=(float*)malloc((size_t)B*H*T*T*sizeof(float));
    float *y=(float*)malloc(tot*sizeof(float));
    if(!p||!y){ free(p); free(y); return -1; }
    attn_score_av(P,q,k,v,B,T,p,y);          /* recompute selection + p */
    attn_score_av_bwd(P,q,k,v,p,dy,B,T,dq,dk,dv);
    free(p); free(y);
    return 0;
}

/* ================================================================
 * FULL RWKV-x MODEL: embedding, head, cross-entropy, from-scratch init, and
 * the whole-model forward/backward assembly (single sequence, B=1).
 * ================================================================ */
#include <stdint.h>

/* ---- Embedding ---- */
void rx_embedding_forward(const float *emb, const int *tokens,
                          int T, int C, float *out)
{
    for (int t=0;t<T;t++)
        memcpy(&out[(size_t)t*C], &emb[(size_t)tokens[t]*C], (size_t)C*sizeof(float));
}

void rx_embedding_backward(const int *tokens, int T, int C, int vocab,
                           const float *d_out, float *d_emb)
{
    memset(d_emb, 0, (size_t)vocab*C*sizeof(float));   /* single-thread scatter-add */
    for (int t=0;t<T;t++){
        float *row=&d_emb[(size_t)tokens[t]*C];
        const float *g=&d_out[(size_t)t*C];
        for (int i=0;i<C;i++) row[i]+=g[i];
    }
}

/* ---- Head / unembedding.  head is [C,vocab] row-major. ---- */
void rx_head_forward(const float *x, const float *head,
                     int T, int C, int vocab, float *logits)
{
    for (int t=0;t<T;t++){
        const float *xt=&x[(size_t)t*C];
        float *lt=&logits[(size_t)t*vocab];
        for (int j=0;j<vocab;j++){
            double s=0.0;                     /* fp64 accumulate: vocab can be large */
            for (int i=0;i<C;i++) s+=(double)xt[i]*(double)head[(size_t)i*vocab+j];
            lt[j]=(float)s;
        }
    }
}

void rx_head_backward(const float *x, const float *head, const float *dlogits,
                      int T, int C, int vocab, float *dx, float *d_head)
{
    memset(d_head, 0, (size_t)C*vocab*sizeof(float));
    memset(dx, 0, (size_t)T*C*sizeof(float));
    for (int t=0;t<T;t++){
        const float *xt=&x[(size_t)t*C];
        const float *dl=&dlogits[(size_t)t*vocab];
        float *dxt=&dx[(size_t)t*C];
        for (int i=0;i<C;i++){
            const float *hrow=&head[(size_t)i*vocab];
            float *dhrow=&d_head[(size_t)i*vocab];
            double acc=0.0; float xi=xt[i];
            for (int j=0;j<vocab;j++){ acc+=(double)dl[j]*(double)hrow[j]; dhrow[j]+=xi*dl[j]; }
            dxt[i]=(float)acc;
        }
    }
}

/* ---- Fused log-softmax cross-entropy ---- */
double rx_cross_entropy_forward(const float *logits, const int *targets,
                                int T, int vocab, int ignore_index)
{
    double loss=0.0; int nv=0;
    for (int t=0;t<T;t++){
        int tg=targets[t];
        if (ignore_index>=0 && tg==ignore_index) continue;
        const float *lt=&logits[(size_t)t*vocab];
        double mx=-1e300; for (int j=0;j<vocab;j++) if (lt[j]>mx) mx=lt[j];
        double se=0.0; for (int j=0;j<vocab;j++) se+=exp((double)lt[j]-mx);
        double lse=mx+log(se);
        loss += lse - (double)lt[tg];
        nv++;
    }
    return nv>0 ? loss/(double)nv : 0.0;
}

double rx_cross_entropy_backward(const float *logits, const int *targets,
                                 int T, int vocab, int ignore_index,
                                 float *dlogits)
{
    /* count valid first (mean normalisation) */
    int nv=0;
    for (int t=0;t<T;t++){ int tg=targets[t];
        if (!(ignore_index>=0 && tg==ignore_index)) nv++; }
    double inv = nv>0 ? 1.0/(double)nv : 0.0;
    double loss=0.0;
    for (int t=0;t<T;t++){
        int tg=targets[t];
        float *dl=&dlogits[(size_t)t*vocab];
        const float *lt=&logits[(size_t)t*vocab];
        if (ignore_index>=0 && tg==ignore_index){
            for (int j=0;j<vocab;j++) dl[j]=0.0f;
            continue;
        }
        double mx=-1e300; for (int j=0;j<vocab;j++) if (lt[j]>mx) mx=lt[j];
        double se=0.0; for (int j=0;j<vocab;j++) se+=exp((double)lt[j]-mx);
        double lse=mx+log(se);
        loss += lse - (double)lt[tg];
        for (int j=0;j<vocab;j++){
            double sm=exp((double)lt[j]-lse);          /* softmax */
            double g=sm - (j==tg?1.0:0.0);
            dl[j]=(float)(g*inv);
        }
    }
    return nv>0 ? loss/(double)nv : 0.0;
}

/* ---- Deterministic PRNG (splitmix64) -> uniform(-0.5, 0.5) ---- */
typedef struct { uint64_t s; } rx_rng;
static inline uint64_t rx_rng_next(rx_rng *r){
    uint64_t z = (r->s += 0x9E3779B97F4A7C15ULL);
    z = (z ^ (z >> 30)) * 0xBF58476D1CE4E5B9ULL;
    z = (z ^ (z >> 27)) * 0x94D049BB133111EBULL;
    return z ^ (z >> 31);
}
static inline float rx_uni(rx_rng *r){   /* uniform in (-0.5, 0.5) */
    return (float)((rx_rng_next(r) >> 11) * (1.0/9007199254740992.0)) - 0.5f;
}

/* ---- Model: registry of owned param tensors + views for the kernels ---- */
typedef struct { char name[48]; float *data; float *grad; int n; } rx_pt_param;

struct rx_ptmodel {
    rx_pt_config pc;
    int C, vocab, n_rwkv, n_attn, n_blocks;
    int *btype;                 /* [n_blocks] 0=rwkv 1=attn */
    int *bidx;                  /* [n_blocks] index within array */
    rx_pt_param *P; int nP, capP;
    float *emb, *ln_out_w, *ln_out_b, *head;
    rx_block_params *rwkv_p;  rx_block_grads *rwkv_g;   /* [n_rwkv] */
    rx_attn_params  *attn_p;  rx_attn_grads  *attn_g;   /* [n_attn] */
    rx_rng rng;
};

/* init-scale codes */
enum { IS_ZEROC, IS_PROJ, IS_LERP, IS_GAMMA, IS_BETA, IS_EMB, IS_HEAD };

static float* rx_reg(rx_ptmodel *m, const char *name, int n, int scale){
    rx_pt_param *p=&m->P[m->nP++];
    snprintf(p->name,sizeof(p->name),"%s",name);
    p->n=n;
    p->data=(float*)malloc((size_t)n*sizeof(float));
    p->grad=(float*)calloc((size_t)n,sizeof(float));
    for (int i=0;i<n;i++){
        float u=rx_uni(&m->rng);
        switch(scale){
            case IS_PROJ:  p->data[i]=u*0.3f; break;
            case IS_LERP:  p->data[i]=u*0.5f; break;
            case IS_GAMMA: p->data[i]=1.0f+0.2f*u; break;
            case IS_BETA:  p->data[i]=0.2f*u; break;
            case IS_EMB:   p->data[i]=u*0.5f; break;
            case IS_HEAD:  p->data[i]=u*0.3f; break;
            default:       p->data[i]=0.0f; break;
        }
    }
    return p->data;
}

void rx_ptmodel_free(rx_ptmodel *m){
    if (!m) return;
    for (int i=0;i<m->nP;i++){ free(m->P[i].data); free(m->P[i].grad); }
    free(m->P); free(m->btype); free(m->bidx);
    free(m->rwkv_p); free(m->rwkv_g); free(m->attn_p); free(m->attn_g);
    free(m);
}

rx_ptmodel* rx_model_new_random(const rx_pt_config *pc, uint64_t seed){
    if (!pc) return NULL;
    int C=pc->n_embd, H=pc->n_head, N=pc->head_size;
    if (H*N!=C || C<=0 || pc->vocab_size<=0 || pc->n_rwkv_layer<=0) return NULL;
    int Dw=pc->Dw,Da=pc->Da,Dg=pc->Dg,Dv=pc->Dv,F=pc->ffn_dim;

    rx_ptmodel *m=(rx_ptmodel*)calloc(1,sizeof(rx_ptmodel));
    m->pc=*pc; m->C=C; m->vocab=pc->vocab_size;
    m->n_rwkv=pc->n_rwkv_layer; m->n_attn=pc->n_attn_layer;
    m->n_blocks=m->n_rwkv+m->n_attn;
    m->rng.s=seed;

    /* upper bound on #params: globals(4) + rwkv(31 each) + attn(11 each) */
    /* upper bound: rwkv layer_id>0 registers 33 params (30 at layer 0),
     * attn 11, globals 4.  The old 31/rwkv bound overflowed m->P for
     * n_rwkv>=4 (heap corruption surfacing at rx_ptmodel_free); use safe 34/12. */
    m->capP = 8 + 34*m->n_rwkv + 12*m->n_attn;
    m->P=(rx_pt_param*)calloc((size_t)m->capP,sizeof(rx_pt_param));
    m->nP=0;

    /* interleave order (mirror rwkvx.c loader) */
    m->btype=(int*)calloc(m->n_blocks,sizeof(int));
    m->bidx =(int*)calloc(m->n_blocks,sizeof(int));
    if (m->n_attn==0){
        for (int i=0;i<m->n_rwkv;i++){ m->btype[i]=0; m->bidx[i]=i; }
    } else {
        int interval=m->n_rwkv/m->n_attn; int pos=0;
        for (int a=0;a<m->n_attn;a++){
            int start=a*interval, end=(a+1)*interval;
            for (int j=start;j<end;j++){ m->btype[pos]=0; m->bidx[pos]=j; pos++; }
            m->btype[pos]=1; m->bidx[pos]=a; pos++;
        }
        /* any RWKV layers beyond n_attn*interval trail at the end */
        for (int j=m->n_attn*interval;j<m->n_rwkv;j++){ m->btype[pos]=0; m->bidx[pos]=j; pos++; }
    }

    /* globals */
    m->emb      = rx_reg(m,"emb",       m->vocab*C, IS_EMB);
    m->rwkv_p=(rx_block_params*)calloc(m->n_rwkv,sizeof(rx_block_params));
    m->rwkv_g=(rx_block_grads *)calloc(m->n_rwkv,sizeof(rx_block_grads));
    m->attn_p=(rx_attn_params *)calloc(m->n_attn>0?m->n_attn:1,sizeof(rx_attn_params));
    m->attn_g=(rx_attn_grads  *)calloc(m->n_attn>0?m->n_attn:1,sizeof(rx_attn_grads));

    char nm[48];
    #define RG(fmt,idx,n,sc) (snprintf(nm,sizeof nm,fmt,idx), rx_reg(m,nm,(n),(sc)))
    /* ---- RWKV-7 blocks ---- */
    for (int i=0;i<m->n_rwkv;i++){
        rx_block_params *P=&m->rwkv_p[i]; rx_block_grads *G=&m->rwkv_g[i];
        rx_tm_params *tm=&P->tm; rx_tm_grads *tg=&G->tm;
        tm->c=C; tm->n_head=H; tm->head_size=N;
        tm->Dw=Dw; tm->Da=Da; tm->Dg=Dg; tm->Dv=Dv; tm->layer_id=i;
        #define TMW(field,cnt,sc) do{ float *d=RG("rwkv%d." #field,i,(cnt),sc); \
            tm->field=d; tg->field=m->P[m->nP-1].grad; }while(0)
        TMW(x_r,C,IS_LERP);TMW(x_k,C,IS_LERP);TMW(x_v,C,IS_LERP);
        TMW(x_w,C,IS_LERP);TMW(x_a,C,IS_LERP);TMW(x_g,C,IS_LERP);
        TMW(receptance_w,C*C,IS_PROJ);TMW(key_w,C*C,IS_PROJ);
        TMW(value_w,C*C,IS_PROJ);TMW(output_w,C*C,IS_PROJ);
        TMW(w0,C,IS_LERP);TMW(w1,C*Dw,IS_PROJ);TMW(w2,Dw*C,IS_PROJ);
        TMW(a0,C,IS_LERP);TMW(a1,C*Da,IS_PROJ);TMW(a2,Da*C,IS_PROJ);
        TMW(g1,C*Dg,IS_PROJ);TMW(g2,Dg*C,IS_PROJ);
        if (i>0){ TMW(v0,C,IS_LERP);TMW(v1,C*Dv,IS_PROJ);TMW(v2,Dv*C,IS_PROJ); }
        else { tm->v0=tm->v1=tm->v2=NULL; tg->v0=tg->v1=tg->v2=NULL; }
        TMW(k_k,C,IS_LERP);TMW(k_a,C,IS_LERP);TMW(r_k,C,IS_LERP);
        TMW(ln_x_w,C,IS_GAMMA);TMW(ln_x_b,C,IS_BETA);
        #undef TMW
        P->cm.c=C; P->cm.ffn_dim=F;
        { float *d;
          d=RG("rwkv%d.ffn_x_k",i,C,IS_LERP);    P->cm.ffn_x_k=d;    G->cm.ffn_x_k=m->P[m->nP-1].grad;
          d=RG("rwkv%d.ffn_key_w",i,C*F,IS_PROJ);P->cm.ffn_key_w=d;  G->cm.ffn_key_w=m->P[m->nP-1].grad;
          d=RG("rwkv%d.ffn_value_w",i,F*C,IS_PROJ);P->cm.ffn_value_w=d;G->cm.ffn_value_w=m->P[m->nP-1].grad;
          d=RG("rwkv%d.ln1_w",i,C,IS_GAMMA); P->ln1_w=d; G->ln1_w=m->P[m->nP-1].grad;
          d=RG("rwkv%d.ln1_b",i,C,IS_BETA);  P->ln1_b=d; G->ln1_b=m->P[m->nP-1].grad;
          d=RG("rwkv%d.ln2_w",i,C,IS_GAMMA); P->ln2_w=d; G->ln2_w=m->P[m->nP-1].grad;
          d=RG("rwkv%d.ln2_b",i,C,IS_BETA);  P->ln2_b=d; G->ln2_b=m->P[m->nP-1].grad;
        }
    }
    /* ---- Sparse-attention blocks ---- */
    for (int i=0;i<m->n_attn;i++){
        rx_attn_params *P=&m->attn_p[i]; rx_attn_grads *G=&m->attn_g[i];
        P->c=C; P->n_head=H; P->head_size=N;
        P->chunk_size=pc->chunk_size<1?1:pc->chunk_size;
        P->moba_top_k=pc->moba_top_k;
        float *d;
        d=RG("attn%d.ln1_w",i,C,IS_GAMMA); P->ln1_w=d; G->ln1_w=m->P[m->nP-1].grad;
        d=RG("attn%d.ln1_b",i,C,IS_BETA);  P->ln1_b=d; G->ln1_b=m->P[m->nP-1].grad;
        d=RG("attn%d.receptance_w",i,C*C,IS_PROJ);P->receptance_w=d;G->receptance_w=m->P[m->nP-1].grad;
        d=RG("attn%d.key_w",i,C*C,IS_PROJ);       P->key_w=d;       G->key_w=m->P[m->nP-1].grad;
        d=RG("attn%d.value_w",i,C*C,IS_PROJ);     P->value_w=d;     G->value_w=m->P[m->nP-1].grad;
        d=RG("attn%d.output_w",i,C*C,IS_PROJ);    P->output_w=d;    G->output_w=m->P[m->nP-1].grad;
        d=RG("attn%d.ln2_w",i,C,IS_GAMMA); P->ln2_w=d; G->ln2_w=m->P[m->nP-1].grad;
        d=RG("attn%d.ln2_b",i,C,IS_BETA);  P->ln2_b=d; G->ln2_b=m->P[m->nP-1].grad;
        P->cm.c=C; P->cm.ffn_dim=F;
        d=RG("attn%d.ffn_x_k",i,C,IS_LERP);        P->cm.ffn_x_k=d;    G->cm.ffn_x_k=m->P[m->nP-1].grad;
        d=RG("attn%d.ffn_key_w",i,C*F,IS_PROJ);    P->cm.ffn_key_w=d;  G->cm.ffn_key_w=m->P[m->nP-1].grad;
        d=RG("attn%d.ffn_value_w",i,F*C,IS_PROJ);  P->cm.ffn_value_w=d;G->cm.ffn_value_w=m->P[m->nP-1].grad;
    }
    #undef RG

    m->ln_out_w = rx_reg(m,"ln_out_w", C, IS_GAMMA);
    m->ln_out_b = rx_reg(m,"ln_out_b", C, IS_BETA);
    m->head     = rx_reg(m,"head",     C*m->vocab, IS_HEAD);
    return m;
}

int  rx_ptmodel_num_params(const rx_ptmodel *m){ return m->nP; }
rx_pt_config rx_ptmodel_config(const rx_ptmodel *m){ return m->pc; }
int  rx_ptmodel_param(rx_ptmodel *m, int i, const char **name,
                      float **data, float **grad, int *n){
    if (i<0||i>=m->nP) return -1;
    if (name) *name=m->P[i].name;
    if (data) *data=m->P[i].data;
    if (grad) *grad=m->P[i].grad;
    if (n)    *n=m->P[i].n;
    return 0;
}
void rx_ptmodel_zero_grad(rx_ptmodel *m){
    for (int i=0;i<m->nP;i++) memset(m->P[i].grad,0,(size_t)m->P[i].n*sizeof(float));
}
const rx_block_params* rx_ptmodel_rwkv(rx_ptmodel *m, int idx){
    return (idx>=0&&idx<m->n_rwkv)?&m->rwkv_p[idx]:NULL;
}
const rx_attn_params* rx_ptmodel_attn(rx_ptmodel *m, int idx){
    return (idx>=0&&idx<m->n_attn)?&m->attn_p[idx]:NULL;
}

/* ---- Activation cache ---- */
rx_pt_acts* rx_pt_acts_new(const rx_ptmodel *m, int T){
    rx_pt_acts *a=(rx_pt_acts*)calloc(1,sizeof(rx_pt_acts));
    a->T=T; a->C=m->C; a->vocab=m->vocab; a->n_blocks=m->n_blocks;
    size_t TC=(size_t)T*m->C;
    a->tokens  =(int*)  calloc(T,sizeof(int));
    a->block_in=(float*)calloc((size_t)(m->n_blocks+1)*TC,sizeof(float));
    a->v_first =(float*)calloc(TC,sizeof(float));
    a->lnout   =(float*)calloc(TC,sizeof(float));
    a->logits  =(float*)calloc((size_t)T*m->vocab,sizeof(float));
    return a;
}
void rx_pt_acts_free(rx_pt_acts *a){
    if (!a) return;
    free(a->tokens); free(a->block_in); free(a->v_first);
    free(a->lnout); free(a->logits); free(a);
}

/* ---- Full forward ---- */
int rx_pretrain_forward(rx_ptmodel *m, const int *tokens, int T, rx_pt_acts *acts){
    int C=m->C; size_t TC=(size_t)T*C;
    for (int t=0;t<T;t++) acts->tokens[t]=tokens[t];
    /* emb -> block_in[0] */
    float *x0=&acts->block_in[0];
    rx_embedding_forward(m->emb, tokens, T, C, x0);
    int rc=0;
    for (int pos=0;pos<m->n_blocks;pos++){
        float *xin =&acts->block_in[(size_t)pos*TC];
        float *xout=&acts->block_in[(size_t)(pos+1)*TC];
        if (m->btype[pos]==0){
            rc|=rx_block_forward_seq(&m->rwkv_p[m->bidx[pos]], xin,
                                     NULL, NULL, acts->v_first, xout, 1, T);
        } else {
            rc|=rx_attn_forward_seq(&m->attn_p[m->bidx[pos]], xin, NULL, xout, 1, T);
        }
    }
    /* final residual = block_in[n_blocks]; ln_out -> lnout; head -> logits */
    float *final=&acts->block_in[(size_t)m->n_blocks*TC];
    rc|=rx_layernorm_forward_seq(final, m->ln_out_w, m->ln_out_b, acts->lnout, 1, T, C);
    rx_head_forward(acts->lnout, m->head, T, C, m->vocab, acts->logits);
    return rc?-1:0;
}

/* ---- Full backward ---- */
double rx_pretrain_backward(rx_ptmodel *m, const rx_pt_acts *acts,
                            const int *targets, int T, int ignore_index){
    int C=m->C, V=m->vocab; size_t TC=(size_t)T*C;
    float *dlogits=(float*)malloc((size_t)T*V*sizeof(float));
    float *dlnout =(float*)malloc(TC*sizeof(float));
    float *dfinal =(float*)malloc(TC*sizeof(float));
    float *dcur   =(float*)malloc(TC*sizeof(float));
    float *dprev  =(float*)malloc(TC*sizeof(float));
    float *g_vf   =(float*)calloc(TC,sizeof(float));   /* accumulated grad on v_first */
    float *dvf_loc=(float*)malloc(TC*sizeof(float));

    /* find analytic grad buffers for globals via registry names is overkill;
     * we know the role pointers' grads live in the registry — look them up. */
    float *d_head=NULL,*d_lnout_w=NULL,*d_lnout_b=NULL,*d_emb=NULL;
    for (int i=0;i<m->nP;i++){
        if (m->P[i].data==m->head)     d_head=m->P[i].grad;
        else if (m->P[i].data==m->ln_out_w) d_lnout_w=m->P[i].grad;
        else if (m->P[i].data==m->ln_out_b) d_lnout_b=m->P[i].grad;
        else if (m->P[i].data==m->emb)      d_emb=m->P[i].grad;
    }

    /* CE -> dlogits, loss */
    double loss=rx_cross_entropy_backward(acts->logits, targets, T, V, ignore_index, dlogits);

    /* head: dlogits -> dlnout + d_head */
    rx_head_backward(acts->lnout, m->head, dlogits, T, C, V, dlnout, d_head);

    /* ln_out: dlnout -> dfinal + d_lnout_{w,b} */
    float *final=&acts->block_in[(size_t)m->n_blocks*TC];
    rx_layernorm_backward(final, m->ln_out_w, m->ln_out_b, dlnout,
                          dfinal, d_lnout_w, d_lnout_b, 1, T, C);

    /* reverse through the interleaved blocks */
    memcpy(dcur, dfinal, TC*sizeof(float));
    for (int pos=m->n_blocks-1;pos>=0;pos--){
        float *xin=&acts->block_in[(size_t)pos*TC];
        if (m->btype[pos]==0){
            int idx=m->bidx[pos];
            if (idx>0){
                rx_block_backward(&m->rwkv_p[idx], xin, NULL, NULL,
                                  acts->v_first, dcur, &m->rwkv_g[idx], dprev,
                                  NULL, NULL, dvf_loc, 1, T);
                for (size_t i=0;i<TC;i++) g_vf[i]+=dvf_loc[i];
            } else {
                /* producer (layer 0): inject accumulated downstream v_first grad */
                rx_block_backward(&m->rwkv_p[idx], xin, NULL, NULL,
                                  NULL, dcur, &m->rwkv_g[idx], dprev,
                                  NULL, NULL, g_vf, 1, T);
            }
        } else {
            rx_attn_backward_full(&m->attn_p[m->bidx[pos]], xin, NULL, dcur,
                                  &m->attn_g[m->bidx[pos]], dprev, NULL, 1, T);
        }
        float *tmp=dcur; dcur=dprev; dprev=tmp;
    }

    /* embedding scatter-add: dcur holds grad on emb output (block_in[0]) */
    rx_embedding_backward(acts->tokens, T, C, V, dcur, d_emb);

    free(dlogits);free(dlnout);free(dfinal);free(dcur);free(dprev);free(g_vf);free(dvf_loc);
    return loss;
}