277 lines
12 KiB
C
277 lines
12 KiB
C
/* SPDX-License-Identifier: AGPL-3.0-or-later */
|
|
/* Copyright (C) 2026 David (Corey) Herrera */
|
|
/*
|
|
* test_attn_bwd.c — Gradient-check for the SPARSE (MoBA-style) ATTENTION BLOCK
|
|
* backward with BASE-weight gradients (rx_attn_backward_full), the counterpart
|
|
* of rwkvx.c `attn_block_forward_fast`:
|
|
* 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
|
|
*
|
|
* Method (same discipline as test_block_bwd.c / docs §5-§6):
|
|
* - A DOUBLE-PRECISION mirror of the whole attention block computes the FD
|
|
* reference loss L = sum(cc * out). fp32 forward noise would swamp the
|
|
* small LN grads and false-fail; the reference is fp64. The double mirror
|
|
* accumulates the softmax + LN reductions in double (cancellation-prone).
|
|
* - The analytic backward runs the normal fp32 path (rx_attn_backward_full).
|
|
* - Central finite differences, PER-TENSOR normalised infinity-norm metric
|
|
* (max|g-fd| / max|fd|). Pass: < 1e-3 for EVERY tensor and the block dx.
|
|
*
|
|
* TOP-K SELECTION (non-differentiable): the mirror recomputes the MoBA chunk
|
|
* selection at the perturbed weights exactly as the kernel does at the base
|
|
* weights. The analytic backward treats the base selection as FIXED (keys of
|
|
* unselected chunks carry weight 0 and vanish from the reductions). These
|
|
* agree ONLY while a perturbation does not flip the selection, so the FD test
|
|
* is run in a selection-STABLE regime:
|
|
* - cases "allsel*" set moba_top_k >= #candidate chunks => all chunks always
|
|
* selected => plain causal attention => selection can never flip.
|
|
* - case "sparse" genuinely drops chunks (moba_top_k < #candidates); with a
|
|
* well-separated random init the gate ordering is stable under the tiny
|
|
* central-difference step, so it still passes. This is the case that
|
|
* exercises the true sparse (masked) backward.
|
|
*/
|
|
#include "lib/test.h"
|
|
#include "rwkvx/rx_pretrain.h"
|
|
|
|
#include <math.h>
|
|
#include <stdlib.h>
|
|
#include <string.h>
|
|
|
|
typedef struct {
|
|
int B, T, H, N, C, F; /* C=H*N, F=ffn_dim */
|
|
int chunk_size, moba_top_k; /* MoBA config */
|
|
} cfg;
|
|
|
|
/* Every learnable tensor in the attention block, fp32. */
|
|
typedef struct {
|
|
float *ln1_w,*ln1_b;
|
|
float *receptance_w,*key_w,*value_w,*output_w;
|
|
float *ln2_w,*ln2_b;
|
|
float *ffn_x_k,*ffn_key_w,*ffn_value_w;
|
|
} AP;
|
|
|
|
static float frand(void){ return (float)rand()/(float)RAND_MAX - 0.5f; }
|
|
static float *rbuf(int n, float scale){
|
|
float *p=malloc((size_t)n*sizeof(float));
|
|
for (int i=0;i<n;i++) p[i]=scale*frand();
|
|
return p;
|
|
}
|
|
|
|
static void ap_alloc(AP *m, const cfg *cf){
|
|
int C=cf->C,F=cf->F;
|
|
m->ln1_w=malloc((size_t)C*sizeof(float)); m->ln1_b=rbuf(C,0.2f);
|
|
m->ln2_w=malloc((size_t)C*sizeof(float)); m->ln2_b=rbuf(C,0.2f);
|
|
for (int i=0;i<C;i++){ m->ln1_w[i]=1.0f+0.2f*frand(); m->ln2_w[i]=1.0f+0.2f*frand(); }
|
|
m->receptance_w=rbuf(C*C,0.3f); m->key_w=rbuf(C*C,0.3f);
|
|
m->value_w=rbuf(C*C,0.3f); m->output_w=rbuf(C*C,0.3f);
|
|
m->ffn_x_k=rbuf(C,1.0f);
|
|
m->ffn_key_w=rbuf(C*F,0.3f);
|
|
m->ffn_value_w=rbuf(F*C,0.3f);
|
|
}
|
|
static void ap_free(AP *m){
|
|
float **a=(float**)m; for (size_t i=0;i<sizeof(AP)/sizeof(float*);i++) free(a[i]);
|
|
}
|
|
static void ap_to_params(const AP *m, const cfg *cf, rx_attn_params *P){
|
|
P->c=cf->C; P->n_head=cf->H; P->head_size=cf->N;
|
|
P->chunk_size=cf->chunk_size; P->moba_top_k=cf->moba_top_k;
|
|
P->ln1_w=m->ln1_w; P->ln1_b=m->ln1_b;
|
|
P->receptance_w=m->receptance_w; P->key_w=m->key_w;
|
|
P->value_w=m->value_w; P->output_w=m->output_w;
|
|
P->ln2_w=m->ln2_w; P->ln2_b=m->ln2_b;
|
|
P->cm.c=cf->C; P->cm.ffn_dim=cf->F;
|
|
P->cm.ffn_x_k=m->ffn_x_k; P->cm.ffn_key_w=m->ffn_key_w; P->cm.ffn_value_w=m->ffn_value_w;
|
|
}
|
|
|
|
/* ---------- fp64 reference forward + scalar loss ---------- */
|
|
static void ln_fp64(const double *x, const float *w, const float *b,
|
|
double *out, int C){
|
|
double mean=0.0; for(int i=0;i<C;i++) mean+=x[i]; mean/=C;
|
|
double var=0.0; for(int i=0;i<C;i++){double d=x[i]-mean;var+=d*d;} var/=C;
|
|
double rstd=1.0/sqrt(var+1e-5);
|
|
for(int i=0;i<C;i++) out[i]=(x[i]-mean)*rstd*(double)w[i]+(double)b[i];
|
|
}
|
|
|
|
static double loss_fp64(const AP *m, const cfg *cf, const float *x,
|
|
const float *cmxp, const float *cc)
|
|
{
|
|
int B=cf->B,T=cf->T,H=cf->H,N=cf->N,C=cf->C,F=cf->F;
|
|
int cs=cf->chunk_size<1?1:cf->chunk_size;
|
|
double scale=1.0/sqrt((double)N);
|
|
size_t tot=(size_t)B*T*C;
|
|
double L=0.0;
|
|
double *h1=malloc(tot*sizeof(double));
|
|
double *q =malloc(tot*sizeof(double));
|
|
double *k =malloc(tot*sizeof(double));
|
|
double *v =malloc(tot*sizeof(double));
|
|
double *y =malloc(tot*sizeof(double));
|
|
double *x1=malloc(tot*sizeof(double));
|
|
double *h2=malloc(tot*sizeof(double));
|
|
double *xd=malloc((size_t)C*sizeof(double));
|
|
double *sh=malloc((size_t)C*sizeof(double));
|
|
double *fk=malloc((size_t)F*sizeof(double));
|
|
|
|
/* h1 = ln1(x); q,k,v = h1 @ {Wq,Wk,Wv} */
|
|
for (int b=0;b<B;b++) for (int t=0;t<T;t++){
|
|
size_t base=((size_t)b*T+t)*C;
|
|
for(int i=0;i<C;i++) xd[i]=(double)x[base+i];
|
|
ln_fp64(xd,m->ln1_w,m->ln1_b,&h1[base],C);
|
|
for(int j=0;j<C;j++){ double sq=0,sk=0,sv=0;
|
|
for(int i=0;i<C;i++){ double hv=h1[base+i];
|
|
sq+=hv*m->receptance_w[i*C+j]; sk+=hv*m->key_w[i*C+j]; sv+=hv*m->value_w[i*C+j]; }
|
|
q[base+j]=sq; k[base+j]=sk; v[base+j]=sv; }
|
|
}
|
|
|
|
/* MoBA-sparse causal attention */
|
|
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;
|
|
int local_chunk=tq/cs, ncand=local_chunk;
|
|
unsigned char csel[T>0?T:1];
|
|
int keepall=(cf->moba_top_k<=0 || cf->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 c2=0;c2<ncand;c2++){ double g=0;
|
|
for(int j=0;j<N;j++){ double mkj=0;
|
|
for(int s=c2*cs;s<(c2+1)*cs;s++) mkj+=k[((size_t)b*T+s)*C+off+j];
|
|
mkj/=(double)cs; g+=q[qb+j]*mkj; }
|
|
gate[c2]=g; csel[c2]=0; }
|
|
for(int it=0;it<cf->moba_top_k;it++){ int arg=-1; double best=-1e300;
|
|
for(int c2=0;c2<ncand;c2++) if(!csel[c2]&&gate[c2]>best){best=gate[c2];arg=c2;}
|
|
if(arg<0)break; csel[arg]=1; }
|
|
}
|
|
double sc[T>0?T:1]; double maxs=-1e300;
|
|
for(int s=0;s<=tq;s++){ int c2=s/cs;
|
|
int inc=(c2==local_chunk)||(c2<ncand&&csel[c2]);
|
|
if(!inc){ sc[s]=-1e300; continue; }
|
|
double d=0; for(int j=0;j<N;j++) d+=q[qb+j]*k[((size_t)b*T+s)*C+off+j];
|
|
d*=scale; sc[s]=d; if(d>maxs)maxs=d; }
|
|
double sum=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 j=0;j<N;j++){ double acc=0;
|
|
for(int s=0;s<=tq;s++){ if(sc[s]<=-1e299)continue;
|
|
acc+=(sc[s]/sum)*v[((size_t)b*T+s)*C+off+j]; }
|
|
y[qb+j]=acc; }
|
|
}
|
|
}
|
|
|
|
/* att_out = y @ Wo; x1 = x + att_out; h2 = ln2(x1) */
|
|
for (int b=0;b<B;b++) for (int t=0;t<T;t++){
|
|
size_t base=((size_t)b*T+t)*C;
|
|
for(int j=0;j<C;j++){ double s=0; for(int i=0;i<C;i++) s+=y[base+i]*m->output_w[i*C+j];
|
|
x1[base+j]=(double)x[base+j]+s; }
|
|
ln_fp64(&x1[base],m->ln2_w,m->ln2_b,&h2[base],C);
|
|
}
|
|
|
|
/* channel-mix over h2; out = x1 + cm; L += cc*out */
|
|
for (int b=0;b<B;b++) for (int t=0;t<T;t++){
|
|
size_t base=((size_t)b*T+t)*C;
|
|
const double *xx=&h2[base];
|
|
const double *prev=(t==0)?NULL:&h2[base-C];
|
|
for(int i=0;i<C;i++){ double pv=prev?prev[i]:(cmxp?(double)cmxp[(size_t)b*C+i]:0.0);
|
|
sh[i]=xx[i]+(pv-xx[i])*(double)m->ffn_x_k[i]; }
|
|
for(int mm=0;mm<F;mm++){ double s=0; for(int i=0;i<C;i++) s+=sh[i]*m->ffn_key_w[i*F+mm];
|
|
double rr=s>0.0?s:0.0; fk[mm]=rr*rr; }
|
|
for(int j=0;j<C;j++){ double s=0; for(int mm=0;mm<F;mm++) s+=fk[mm]*m->ffn_value_w[mm*C+j];
|
|
double out=x1[base+j]+s; L+=(double)cc[base+j]*out; }
|
|
}
|
|
|
|
free(h1);free(q);free(k);free(v);free(y);free(x1);free(h2);
|
|
free(xd);free(sh);free(fk);
|
|
return L;
|
|
}
|
|
|
|
/* FD-check an array (nv elems) vs analytic grad g — per-tensor rel-err. */
|
|
static float fd_array(float *arr, const float *g, int nv,
|
|
const AP *m, const cfg *cf, const float *x, const float *cmxp,
|
|
const float *cc, const char *nm, int verbose)
|
|
{
|
|
float maxdiff=0.0f, maxfd=0.0f;
|
|
for (int e=0;e<nv;e++){
|
|
float o=arr[e]; float h=1e-3f*(1.0f+fabsf(o));
|
|
arr[e]=o+h; double Lp=loss_fp64(m,cf,x,cmxp,cc);
|
|
arr[e]=o-h; double Lm=loss_fp64(m,cf,x,cmxp,cc);
|
|
arr[e]=o;
|
|
float fd=(float)((Lp-Lm)/(2.0*(double)h));
|
|
float d=fabsf(g[e]-fd);
|
|
if (d>maxdiff) maxdiff=d;
|
|
if (fabsf(fd)>maxfd) maxfd=fabsf(fd);
|
|
}
|
|
float re=maxdiff/(maxfd+1e-9f);
|
|
if (verbose) printf(" %-14s rel=%.2e (max|g-fd|=%.2e max|fd|=%.2e)\n",
|
|
nm,re,maxdiff,maxfd);
|
|
return re;
|
|
}
|
|
|
|
static int run_case(const char *name, cfg cf, unsigned seed){
|
|
srand(seed);
|
|
cf.C = cf.H*cf.N;
|
|
int C=cf.C, B=cf.B, T=cf.T;
|
|
size_t tot=(size_t)B*T*C;
|
|
|
|
AP m; ap_alloc(&m,&cf);
|
|
rx_attn_params P; ap_to_params(&m,&cf,&P);
|
|
|
|
float *x=rbuf((int)tot,1.0f);
|
|
float *cc=rbuf((int)tot,1.0f);
|
|
float *dy=malloc(tot*sizeof(float));
|
|
for (size_t i=0;i<tot;i++) dy[i]=cc[i]; /* L = sum cc*out => dy = cc */
|
|
|
|
/* analytic grads */
|
|
AP gm; memset(&gm,0,sizeof gm);
|
|
#define GA(f,n) gm.f=calloc((size_t)(n),sizeof(float))
|
|
GA(ln1_w,C);GA(ln1_b,C);GA(ln2_w,C);GA(ln2_b,C);
|
|
GA(receptance_w,C*C);GA(key_w,C*C);GA(value_w,C*C);GA(output_w,C*C);
|
|
GA(ffn_x_k,C);GA(ffn_key_w,C*cf.F);GA(ffn_value_w,cf.F*C);
|
|
#undef GA
|
|
rx_attn_grads G; memset(&G,0,sizeof G);
|
|
G.ln1_w=gm.ln1_w;G.ln1_b=gm.ln1_b;G.ln2_w=gm.ln2_w;G.ln2_b=gm.ln2_b;
|
|
G.receptance_w=gm.receptance_w;G.key_w=gm.key_w;G.value_w=gm.value_w;G.output_w=gm.output_w;
|
|
G.cm.ffn_x_k=gm.ffn_x_k;G.cm.ffn_key_w=gm.ffn_key_w;G.cm.ffn_value_w=gm.ffn_value_w;
|
|
|
|
float *dx=calloc(tot,sizeof(float));
|
|
rx_attn_backward_full(&P,x,NULL,dy,&G,dx,NULL,B,T);
|
|
|
|
printf(" [%s] B=%d T=%d H=%d N=%d C=%d F=%d chunk=%d topk=%d\n",
|
|
name,B,T,cf.H,cf.N,C,cf.F,cf.chunk_size,cf.moba_top_k);
|
|
float wmax=0.0f;
|
|
#define CK(f,n) do{ float e=fd_array(m.f,gm.f,(n),&m,&cf,x,NULL,cc,#f,1); \
|
|
if(e>wmax)wmax=e; }while(0)
|
|
CK(receptance_w,C*C);CK(key_w,C*C);CK(value_w,C*C);CK(output_w,C*C);
|
|
CK(ln1_w,C);CK(ln1_b,C);CK(ln2_w,C);CK(ln2_b,C);
|
|
CK(ffn_x_k,C);CK(ffn_key_w,C*cf.F);CK(ffn_value_w,cf.F*C);
|
|
#undef CK
|
|
{ float e=fd_array(x,dx,(int)tot,&m,&cf,x,NULL,cc,"dx(block in)",1);
|
|
if(e>wmax)wmax=e; }
|
|
printf(" -> worst rel-err = %.2e (%s)\n", wmax, wmax<1e-3f?"PASS":"FAIL");
|
|
|
|
ap_free(&m); ap_free(&gm);
|
|
free(x);free(cc);free(dy);free(dx);
|
|
return wmax<1e-3f;
|
|
}
|
|
|
|
static void test_attn_backward(void){
|
|
printf(" Sparse-attention block backward (base-weight grads) — fp64 FD:\n");
|
|
int ok=1;
|
|
/* all-selected (moba_top_k >= #candidate chunks) => full causal, selection
|
|
* can never flip under the FD step. */
|
|
ok &= run_case("allsel s1", (cfg){.B=1,.T=6,.H=2,.N=4,.F=12,
|
|
.chunk_size=2,.moba_top_k=99}, 11);
|
|
ok &= run_case("allsel bh", (cfg){.B=2,.T=5,.H=3,.N=6,.F=20,
|
|
.chunk_size=2,.moba_top_k=99}, 23);
|
|
ok &= run_case("allsel cs3",(cfg){.B=1,.T=7,.H=2,.N=4,.F=16,
|
|
.chunk_size=3,.moba_top_k=0 /*<=0=all*/}, 37);
|
|
/* genuinely sparse: drops chunks; passes while gate ordering is stable. */
|
|
ok &= run_case("sparse ", (cfg){.B=1,.T=8,.H=2,.N=4,.F=12,
|
|
.chunk_size=2,.moba_top_k=1}, 41);
|
|
ASSERT_TRUE(ok);
|
|
}
|
|
|
|
int main(void){
|
|
printf("\nSparse-attention block backward gradient check\n");
|
|
RUN_TEST(attn_backward);
|
|
TEST_REPORT();
|
|
return __ft_failed>0 ? 1 : 0;
|
|
}
|