/* SPDX-License-Identifier: AGPL-3.0-or-later */ /* Copyright (C) 2026 David (Corey) Herrera */ /* * test_wkv7_bwd.c โ€” Gradient-check for the RWKV-7 WKV7 backward-through-time. * * Two independent gates (docs/wkv7_backward_reference.md ยง6): * * A. Parity vs the verified fp64 PyTorch oracle. The fixed-seed fixtures in * wkv7_fixture.h (generated by gen_wkv7_fixture.py from wkv7_reference.py) * carry inputs, y, and reference grads dr,dw,dk,dv,da,db. We run the C * forward + backward on the SAME inputs and diff. Pass: max rel-err<1e-3 * on every grad (and on y). * * B. Central finite differences in pure C (self-contained, no torch). For a * scalar loss L=sum(c*y) with fixed random c, FD-check every input * r,w,k,v,a,b AND the initial state carry S_{-1} (-> dS) independently. * Pass: max rel-err<1e-3 for all seven. * * Prints per-grad relative error for both gates. */ #include "lib/test.h" #include "rwkvx/rx_pretrain.h" #include "rwkvx/tests/wkv7_fixture.h" #include #include #include /* max relative error: max_e |g[e]-ref[e]| / (|ref[e]| + eps) */ static float rel_err(const float *g, const float *ref, int n) { float worst = 0.0f; for (int e = 0; e < n; e++) { float d = fabsf(g[e] - ref[e]); float r = d / (fabsf(ref[e]) + 1e-6f); if (r > worst) worst = r; } return worst; } /* ================================================================ * Gate A: parity vs oracle fixture * ================================================================ */ static int gate_a_case(const char *name, int B, int T, int H, int N, const float *r, const float *w, const float *k, const float *v, const float *a, const float *b, const float *dy, const float *y_ref, const float *dr_ref, const float *dw_ref, const float *dk_ref, const float *dv_ref, const float *da_ref, const float *db_ref) { int nv = B * T * H * N; float *y = (float *)malloc(nv * sizeof(float)); float *dr = (float *)malloc(nv * sizeof(float)); float *dw = (float *)malloc(nv * sizeof(float)); float *dk = (float *)malloc(nv * sizeof(float)); float *dv = (float *)malloc(nv * sizeof(float)); float *da = (float *)malloc(nv * sizeof(float)); float *db = (float *)malloc(nv * sizeof(float)); rx_wkv7_forward_seq(r, w, k, v, a, b, NULL, y, NULL, B, T, H, N); rx_wkv7_backward(r, w, k, v, a, b, dy, NULL, dr, dw, dk, dv, da, db, NULL, B, T, H, N); float ey = rel_err(y, y_ref, nv); float edr = rel_err(dr, dr_ref, nv); float edw = rel_err(dw, dw_ref, nv); float edk = rel_err(dk, dk_ref, nv); float edv = rel_err(dv, dv_ref, nv); float eda = rel_err(da, da_ref, nv); float edb = rel_err(db, db_ref, nv); printf(" [A %-10s B=%d T=%2d H=%d N=%2d] y=%.2e | dr=%.2e dw=%.2e " "dk=%.2e dv=%.2e da=%.2e db=%.2e\n", name, B, T, H, N, ey, edr, edw, edk, edv, eda, edb); free(y); free(dr); free(dw); free(dk); free(dv); free(da); free(db); int ok = (ey < 1e-3f) && (edr < 1e-3f) && (edw < 1e-3f) && (edk < 1e-3f) && (edv < 1e-3f) && (eda < 1e-3f) && (edb < 1e-3f); return ok; } static void test_gate_a_oracle_parity(void) { printf(" Gate A โ€” parity vs verified fp64 oracle fixture:\n"); int ok = 1; ok &= gate_a_case("canonical", C0_B, C0_T, C0_H, C0_N, c0_r, c0_w, c0_k, c0_v, c0_a, c0_b, c0_dy, c0_y, c0_dr, c0_dw, c0_dk, c0_dv, c0_da, c0_db); ok &= gate_a_case("multihead", C1_B, C1_T, C1_H, C1_N, c1_r, c1_w, c1_k, c1_v, c1_a, c1_b, c1_dy, c1_y, c1_dr, c1_dw, c1_dk, c1_dv, c1_da, c1_db); ok &= gate_a_case("batched", C2_B, C2_T, C2_H, C2_N, c2_r, c2_w, c2_k, c2_v, c2_a, c2_b, c2_dy, c2_y, c2_dr, c2_dw, c2_dk, c2_dv, c2_da, c2_db); ASSERT_TRUE(ok); } /* ================================================================ * Gate B: pure-C central finite differences (incl. dS carry) * ================================================================ */ /* fp64 reference forward, mirroring rx_wkv7_forward_seq exactly, so the FD * reference loss is free of fp32 forward noise (otherwise single-element * perturbations of small inputs like a/b lose precision to cancellation). * Loss L = sum(c*y). */ static double loss_of(const float *r, const float *w, const float *k, const float *v, const float *a, const float *b, const float *s_init, const float *c, int B, int T, int H, int N) { double *S = (double *)malloc((size_t)B * H * N * N * sizeof(double)); for (int e = 0; e < B*H*N*N; e++) S[e] = s_init ? (double)s_init[e] : 0.0; double L = 0.0; 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 = ((bch*H + h)*N + i)*N; double sa = 0.0; for (int j = 0; j < N; j++) sa += (double)a[(((bch*T+t)*H+h)*N)+j] * S[sbase+j]; double vi = (double)v[(((bch*T+t)*H+h)*N)+i]; double yy = 0.0; for (int j = 0; j < N; j++) { int vj = (((bch*T+t)*H+h)*N)+j; S[sbase+j] = S[sbase+j]*(double)w[vj] + (double)k[vj]*vi + sa*(double)b[vj]; yy += S[sbase+j] * (double)r[vj]; } L += (double)c[(((bch*T+t)*H+h)*N)+i] * yy; } free(S); return L; } /* FD-check one input tensor `x` (nv elements) against analytic grad `g`. */ static float fd_check_input(float *x, const float *g, int nv, const float *r, const float *w, const float *k, const float *v, const float *a, const float *b, const float *s_init, const float *c, int B, int T, int H, int N) { float worst = 0.0f; for (int e = 0; e < nv; e++) { float orig = x[e]; float hstep = 1e-3f * (1.0f + fabsf(orig)); x[e] = orig + hstep; double Lp = loss_of(r, w, k, v, a, b, s_init, c, B, T, H, N); x[e] = orig - hstep; double Lm = loss_of(r, w, k, v, a, b, s_init, c, B, T, H, N); x[e] = orig; float fd = (float)((Lp - Lm) / (2.0 * (double)hstep)); float rerr = fabsf(g[e] - fd) / (fabsf(fd) + 1e-6f); if (rerr > worst) worst = rerr; } return worst; } /* FD-check the initial state carry S_{-1} against dS_out (nvs elements). */ static float fd_check_state(float *s_init, const float *dS, int nvs, const float *r, const float *w, const float *k, const float *v, const float *a, const float *b, const float *c, int B, int T, int H, int N) { float worst = 0.0f; for (int e = 0; e < nvs; e++) { float orig = s_init[e]; float hstep = 1e-3f * (1.0f + fabsf(orig)); s_init[e] = orig + hstep; double Lp = loss_of(r, w, k, v, a, b, s_init, c, B, T, H, N); s_init[e] = orig - hstep; double Lm = loss_of(r, w, k, v, a, b, s_init, c, B, T, H, N); s_init[e] = orig; float fd = (float)((Lp - Lm) / (2.0 * (double)hstep)); float rerr = fabsf(dS[e] - fd) / (fabsf(fd) + 1e-6f); if (rerr > worst) worst = rerr; } return worst; } static float frand(void) { return (float)rand() / (float)RAND_MAX - 0.5f; } static void test_gate_b_finite_diff(void) { printf(" Gate B โ€” central finite differences in C (incl. dS carry):\n"); int B = 1, T = 6, H = 2, N = 6; int nv = B * T * H * N, nvs = B * H * N * N; srand(1234); float *r = malloc(nv*4), *w = malloc(nv*4), *k = malloc(nv*4); float *v = malloc(nv*4), *a = malloc(nv*4), *b = malloc(nv*4); float *c = malloc(nv*4), *s_init = malloc(nvs*4); for (int e = 0; e < nv; e++) { r[e] = 0.5f*frand(); k[e] = 0.5f*frand(); v[e] = 0.5f*frand(); a[e] = 0.5f*frand(); b[e] = 0.5f*frand(); c[e] = frand(); /* effective decay in (0,1) */ w[e] = 0.3f + 0.6f*((float)rand()/(float)RAND_MAX); } /* nonzero initial state carry -> genuinely exercises the reverse scan */ for (int e = 0; e < nvs; e++) s_init[e] = 0.3f*frand(); /* analytic grads: dy = c (since L = sum c*y) */ float *dr = malloc(nv*4), *dw = malloc(nv*4), *dk = malloc(nv*4); float *dv = malloc(nv*4), *da = malloc(nv*4), *db = malloc(nv*4); float *dS = malloc(nvs*4); rx_wkv7_backward(r, w, k, v, a, b, c, s_init, dr, dw, dk, dv, da, db, dS, B, T, H, N); float er = fd_check_input(r, dr, nv, r,w,k,v,a,b,s_init,c, B,T,H,N); float ek = fd_check_input(k, dk, nv, r,w,k,v,a,b,s_init,c, B,T,H,N); float ev = fd_check_input(v, dv, nv, r,w,k,v,a,b,s_init,c, B,T,H,N); float ew = fd_check_input(w, dw, nv, r,w,k,v,a,b,s_init,c, B,T,H,N); float eb = fd_check_input(b, db, nv, r,w,k,v,a,b,s_init,c, B,T,H,N); float ea = fd_check_input(a, da, nv, r,w,k,v,a,b,s_init,c, B,T,H,N); float es = fd_check_state(s_init, dS, nvs, r,w,k,v,a,b,c, B,T,H,N); printf(" [B fd B=%d T=%2d H=%d N=%2d] dr=%.2e dk=%.2e dv=%.2e " "dw=%.2e db=%.2e da=%.2e | dS=%.2e\n", B, T, H, N, er, ek, ev, ew, eb, ea, es); int ok = (er<1e-3f)&&(ek<1e-3f)&&(ev<1e-3f)&&(ew<1e-3f) &&(eb<1e-3f)&&(ea<1e-3f)&&(es<1e-3f); free(r);free(w);free(k);free(v);free(a);free(b);free(c);free(s_init); free(dr);free(dw);free(dk);free(dv);free(da);free(db);free(dS); ASSERT_TRUE(ok); } int main(void) { printf("\nWKV7 backward-through-time gradient check\n"); RUN_TEST(gate_a_oracle_parity); RUN_TEST(gate_b_finite_diff); TEST_REPORT(); return __ft_failed > 0 ? 1 : 0; }