Foundry/cuda_kernels/tests/test_timemix_cuda.cu

/* SPDX-License-Identifier: AGPL-3.0-or-later */
/* Copyright (C) 2026 David (Corey) Herrera */
/*
 * test_timemix_cuda.cu — parity harness for the GPU RWKV-7 time-mix block
 * (cuda_kernels/timemix.cu) vs the CPU oracle rx_timemix_forward_seq /
 * rx_timemix_backward (rwkvx/rx_pretrain.c), on GPU 0.  Fixed-seed generated
 * fixtures at the real config incl. N=64.  Parity metric: max-abs-error
 * relative to tensor scale (inf-norm), bar < 1e-3 (same as the WKV7 harness;
 * immune to signed-reduction near-zero entries).  Covers the block output +
 * dx, both token-shift carry grads, and EVERY learnable time-mix param grad.
 */
#include <cuda_runtime.h>
#include <cublas_v2.h>
#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <string.h>
#include <math.h>
#include <time.h>

#include "cuda_kernels/timemix.cu"

extern "C" int rx_timemix_forward_seq(const rx_tm_params*,const float*,const float*,float*,float*,int,int);
extern "C" int rx_timemix_backward(const rx_tm_params*,const float*,const float*,const float*,
    const float*,rx_tm_grads*,float*,float*,float*,int,int);

static uint64_t sm_state;
static uint64_t sm_next(void){ uint64_t z=(sm_state+=0x9E3779B97F4A7C15ULL);
    z=(z^(z>>30))*0xBF58476D1CE4E5B9ULL; z=(z^(z>>27))*0x94D049BB133111EBULL; return z^(z>>31); }
static float sm_uf(void){ return (float)((sm_next()>>11)*(1.0/9007199254740992.0)); }
static float sm_sym(float s){ return (2.0f*sm_uf()-1.0f)*s; }
static float* rnd(size_t n,float s){ float*p=(float*)malloc(n*4); for(size_t i=0;i<n;i++)p[i]=sm_sym(s); return p; }
static float* rndln(size_t n){ float*p=(float*)malloc(n*4); for(size_t i=0;i<n;i++)p[i]=1.0f+0.2f*sm_sym(1.0f); return p; }
static float* zeros(size_t n){ return (float*)calloc(n,4); }

static int g_fail=0;
static void chk(const char*tname,const float*gpu,const float*cpu,size_t n){
    if(!gpu||!cpu) return;
    double maxabs=0.0,scale=0.0;
    for(size_t i=0;i<n;i++){ double d=fabs((double)gpu[i]-(double)cpu[i]); if(d>maxabs)maxabs=d;
        double c=fabs((double)cpu[i]); if(c>scale)scale=c; }
    double rel=maxabs/(scale+1e-30); const char*st=(rel<1e-3)?"PASS":"FAIL"; if(rel>=1e-3)g_fail=1;
    printf("    %-14s rel=%.3e maxabs=%.3e scale=%.3e [%s]\n",tname,rel,maxabs,scale,st);
}

typedef struct { int B,T,H,N,C,Dw,Da,Dg,Dv,layer_id; } cfg;

/* allocate a full grads struct (all buffers) */
static void grads_alloc(rx_tm_grads*G,const cfg*cf){
    int C=cf->C,Dw=cf->Dw,Da=cf->Da,Dg=cf->Dg,Dv=cf->Dv;
    G->x_r=zeros(C);G->x_k=zeros(C);G->x_v=zeros(C);G->x_w=zeros(C);G->x_a=zeros(C);G->x_g=zeros(C);
    G->receptance_w=zeros((size_t)C*C);G->key_w=zeros((size_t)C*C);
    G->value_w=zeros((size_t)C*C);G->output_w=zeros((size_t)C*C);
    G->w0=zeros(C);G->w1=zeros((size_t)C*Dw);G->w2=zeros((size_t)Dw*C);
    G->a0=zeros(C);G->a1=zeros((size_t)C*Da);G->a2=zeros((size_t)Da*C);
    G->g1=zeros((size_t)C*Dg);G->g2=zeros((size_t)Dg*C);
    if(cf->layer_id>0){G->v0=zeros(C);G->v1=zeros((size_t)C*Dv);G->v2=zeros((size_t)Dv*C);}
    else {G->v0=G->v1=G->v2=NULL;}
    G->k_k=zeros(C);G->k_a=zeros(C);G->r_k=zeros(C);G->ln_x_w=zeros(C);G->ln_x_b=zeros(C);
}

static void cmp_grads(rx_tm_grads*a,rx_tm_grads*b,const cfg*cf){
    int C=cf->C,Dw=cf->Dw,Da=cf->Da,Dg=cf->Dg,Dv=cf->Dv;
    chk("x_r",a->x_r,b->x_r,C);chk("x_k",a->x_k,b->x_k,C);chk("x_v",a->x_v,b->x_v,C);
    chk("x_w",a->x_w,b->x_w,C);chk("x_a",a->x_a,b->x_a,C);chk("x_g",a->x_g,b->x_g,C);
    chk("receptance_w",a->receptance_w,b->receptance_w,(size_t)C*C);
    chk("key_w",a->key_w,b->key_w,(size_t)C*C);
    chk("value_w",a->value_w,b->value_w,(size_t)C*C);
    chk("output_w",a->output_w,b->output_w,(size_t)C*C);
    chk("w0",a->w0,b->w0,C);chk("w1",a->w1,b->w1,(size_t)C*Dw);chk("w2",a->w2,b->w2,(size_t)Dw*C);
    chk("a0",a->a0,b->a0,C);chk("a1",a->a1,b->a1,(size_t)C*Da);chk("a2",a->a2,b->a2,(size_t)Da*C);
    chk("g1",a->g1,b->g1,(size_t)C*Dg);chk("g2",a->g2,b->g2,(size_t)Dg*C);
    if(cf->layer_id>0){chk("v0",a->v0,b->v0,C);chk("v1",a->v1,b->v1,(size_t)C*Dv);chk("v2",a->v2,b->v2,(size_t)Dv*C);}
    chk("k_k",a->k_k,b->k_k,C);chk("k_a",a->k_a,b->k_a,C);chk("r_k",a->r_k,b->r_k,C);
    chk("ln_x_w",a->ln_x_w,b->ln_x_w,C);chk("ln_x_b",a->ln_x_b,b->ln_x_b,C);
}

static void run_case(const char*name, cfg cf, uint64_t seed){
    int C=cf.C,Dw=cf.Dw,Da=cf.Da,Dg=cf.Dg,Dv=cf.Dv; size_t BT=(size_t)cf.B*cf.T, tot=BT*C;
    sm_state=seed;
    rx_tm_params P; memset(&P,0,sizeof(P));
    P.c=C;P.n_head=cf.H;P.head_size=cf.N;P.Dw=Dw;P.Da=Da;P.Dg=Dg;P.Dv=Dv;P.layer_id=cf.layer_id;
    P.x_r=rnd(C,0.5f);P.x_k=rnd(C,0.5f);P.x_v=rnd(C,0.5f);P.x_w=rnd(C,0.5f);P.x_a=rnd(C,0.5f);P.x_g=rnd(C,0.5f);
    P.receptance_w=rnd((size_t)C*C,0.3f);P.key_w=rnd((size_t)C*C,0.3f);
    P.value_w=rnd((size_t)C*C,0.3f);P.output_w=rnd((size_t)C*C,0.3f);
    P.w0=rnd(C,0.5f);P.w1=rnd((size_t)C*Dw,0.3f);P.w2=rnd((size_t)Dw*C,0.3f);
    P.a0=rnd(C,0.5f);P.a1=rnd((size_t)C*Da,0.3f);P.a2=rnd((size_t)Da*C,0.3f);
    P.g1=rnd((size_t)C*Dg,0.3f);P.g2=rnd((size_t)Dg*C,0.3f);
    if(cf.layer_id>0){P.v0=rnd(C,0.5f);P.v1=rnd((size_t)C*Dv,0.3f);P.v2=rnd((size_t)Dv*C,0.3f);}
    P.k_k=rnd(C,0.5f);P.k_a=rnd(C,0.5f);P.r_k=rnd(C,0.5f);P.ln_x_w=rndln(C);P.ln_x_b=rnd(C,0.2f);

    float*x=rnd(tot,1.0f); float*xprev=rnd((size_t)cf.B*C,1.0f);
    float*vf_c=rnd(tot,1.0f); float*vf_g=(float*)malloc(tot*4); memcpy(vf_g,vf_c,tot*4);
    float*out_c=(float*)malloc(tot*4); float*out_g=(float*)malloc(tot*4);
    float*dy=rnd(tot,1.0f);

    /* forward */
    rx_timemix_forward_seq(&P,x,xprev,vf_c,out_c,cf.B,cf.T);
    if(timemix_forward_cuda(&P,x,xprev,vf_g,out_g,cf.B,cf.T)){printf("  %s GPU fwd ERR\n",name);g_fail=1;return;}

    /* backward */
    rx_tm_grads Gc,Gg; grads_alloc(&Gc,&cf); grads_alloc(&Gg,&cf);
    float*dx_c=zeros(tot),*dx_g=zeros(tot);
    float*dxp_c=zeros((size_t)cf.B*C),*dxp_g=zeros((size_t)cf.B*C);
    /* d_v_first: OUTPUT for layer>0; for layer0 an optional INPUT (downstream grad) */
    float*dvf_out_c=NULL,*dvf_out_g=NULL,*dvf_in=NULL;
    if(cf.layer_id>0){ dvf_out_c=zeros(tot); dvf_out_g=zeros(tot); }
    else { dvf_in=rnd(tot,1.0f); }  /* layer0: downstream grad on produced v_first */

    rx_timemix_backward(&P,x,xprev,vf_c,dy,&Gc,dx_c,dxp_c,
                        cf.layer_id>0?dvf_out_c:dvf_in,cf.B,cf.T);
    if(timemix_backward_cuda(&P,x,xprev,vf_g,dy,&Gg,dx_g,dxp_g,
                        cf.layer_id>0?dvf_out_g:dvf_in,cf.B,cf.T)){printf("  %s GPU bwd ERR\n",name);g_fail=1;return;}

    printf("  %s (B=%d T=%d H=%d N=%d C=%d layer=%d)\n",name,cf.B,cf.T,cf.H,cf.N,C,cf.layer_id);
    chk("out",out_g,out_c,tot);
    if(cf.layer_id==0) chk("v_first",vf_g,vf_c,tot);
    chk("dx",dx_g,dx_c,tot);
    chk("d_xprev",dxp_g,dxp_c,(size_t)cf.B*C);
    if(cf.layer_id>0) chk("d_v_first",dvf_out_g,dvf_out_c,tot);
    cmp_grads(&Gg,&Gc,&cf);
}

static void bench(cfg cf,int iters){
    int C=cf.C,Dw=cf.Dw,Da=cf.Da,Dg=cf.Dg,Dv=cf.Dv; size_t BT=(size_t)cf.B*cf.T, tot=BT*C;
    sm_state=99;
    rx_tm_params P; memset(&P,0,sizeof(P));
    P.c=C;P.n_head=cf.H;P.head_size=cf.N;P.Dw=Dw;P.Da=Da;P.Dg=Dg;P.Dv=Dv;P.layer_id=cf.layer_id;
    P.x_r=rnd(C,0.5f);P.x_k=rnd(C,0.5f);P.x_v=rnd(C,0.5f);P.x_w=rnd(C,0.5f);P.x_a=rnd(C,0.5f);P.x_g=rnd(C,0.5f);
    P.receptance_w=rnd((size_t)C*C,0.3f);P.key_w=rnd((size_t)C*C,0.3f);
    P.value_w=rnd((size_t)C*C,0.3f);P.output_w=rnd((size_t)C*C,0.3f);
    P.w0=rnd(C,0.5f);P.w1=rnd((size_t)C*Dw,0.3f);P.w2=rnd((size_t)Dw*C,0.3f);
    P.a0=rnd(C,0.5f);P.a1=rnd((size_t)C*Da,0.3f);P.a2=rnd((size_t)Da*C,0.3f);
    P.g1=rnd((size_t)C*Dg,0.3f);P.g2=rnd((size_t)Dg*C,0.3f);
    if(cf.layer_id>0){P.v0=rnd(C,0.5f);P.v1=rnd((size_t)C*Dv,0.3f);P.v2=rnd((size_t)Dv*C,0.3f);}
    P.k_k=rnd(C,0.5f);P.k_a=rnd(C,0.5f);P.r_k=rnd(C,0.5f);P.ln_x_w=rndln(C);P.ln_x_b=rnd(C,0.2f);
    float*x=rnd(tot,1.0f),*xprev=rnd((size_t)cf.B*C,1.0f),*vf=rnd(tot,1.0f),*out=(float*)malloc(tot*4);
    timemix_forward_cuda(&P,x,xprev,vf,out,cf.B,cf.T); /* warmup */
    struct timespec t0,t1; clock_gettime(CLOCK_MONOTONIC,&t0);
    for(int i=0;i<iters;i++) timemix_forward_cuda(&P,x,xprev,vf,out,cf.B,cf.T);
    clock_gettime(CLOCK_MONOTONIC,&t1);
    double sec=(t1.tv_sec-t0.tv_sec)+(t1.tv_nsec-t0.tv_nsec)*1e-9;
    double toks=(double)cf.B*cf.T*iters;
    printf("  timemix fwd B=%d T=%d H=%d N=%d C=%d layer=%d: %d it in %.4fs -> %.0f tokens/sec (one K80 die; incl H2D+alloc)\n",
        cf.B,cf.T,cf.H,cf.N,C,cf.layer_id,iters,sec,toks/sec);
}

int main(void){
    int dev=0; cudaSetDevice(dev); cudaDeviceProp p; cudaGetDeviceProperties(&p,dev);
    printf("=== GPU time-mix FORWARD+BACKWARD parity vs CPU oracle ===\n");
    printf("device %d: %s cc %d.%d\n\n",dev,p.name,p.major,p.minor);
    run_case("small L0", (cfg){.B=1,.T=5,.H=2,.N=4,.C=8,.Dw=3,.Da=3,.Dg=3,.Dv=0,.layer_id=0},11);
    run_case("small L1", (cfg){.B=1,.T=6,.H=2,.N=4,.C=8,.Dw=3,.Da=3,.Dg=3,.Dv=3,.layer_id=1},23);
    run_case("N64 L1",   (cfg){.B=2,.T=16,.H=2,.N=64,.C=128,.Dw=32,.Da=32,.Dg=32,.Dv=32,.layer_id=1},37);
    run_case("N64 L0",   (cfg){.B=1,.T=16,.H=2,.N=64,.C=128,.Dw=32,.Da=32,.Dg=32,.Dv=0,.layer_id=0},41);
    printf("\n=== throughput (time-mix forward, one die) ===\n");
    bench((cfg){.B=8,.T=512,.H=12,.N=64,.C=768,.Dw=64,.Da=64,.Dg=64,.Dv=64,.layer_id=1},20);
    printf("\n%s\n", g_fail?"RESULT: FAIL":"RESULT: PASS");
    return g_fail;
}