David
/
Foundry
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
Foundry/rwkvx/rx_pretrain.c

1665 lines
73 KiB
C
Raw Blame History

/* 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;
}
Reference in New Issue View Git Blame Copy Permalink