97 lines
4.7 KiB
Plaintext
97 lines
4.7 KiB
Plaintext
/* SPDX-License-Identifier: AGPL-3.0-or-later */
|
|
/* Copyright (C) 2026 David (Corey) Herrera */
|
|
/*
|
|
* test_channelmix_cuda.cu — parity harness for GPU RWKV-7 channel-mix
|
|
* (cuda_kernels/channelmix.cu) vs the CPU oracle rx_channelmix_forward_seq /
|
|
* rx_channelmix_backward (rwkvx/rx_pretrain.c), on GPU 0. Fixed-seed fixtures
|
|
* incl. the real config C=768, F=4*C=3072. Parity metric: max-abs-error /
|
|
* tensor-scale (inf-norm), bar < 1e-3. Covers out, dx, d_xprev_init, and every
|
|
* param grad (ffn_x_k, ffn_key_w, ffn_value_w).
|
|
*/
|
|
#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/channelmix.cu"
|
|
|
|
extern "C" int rx_channelmix_forward_seq(const rx_cm_params*,const float*,const float*,float*,int,int);
|
|
extern "C" int rx_channelmix_backward(const rx_cm_params*,const float*,const float*,const float*,
|
|
rx_cm_grads*,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* 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(" %-13s rel=%.3e maxabs=%.3e scale=%.3e [%s]\n",tname,rel,maxabs,scale,st);
|
|
}
|
|
|
|
static void mkP(rx_cm_params*P,int C,int F){
|
|
P->c=C; P->ffn_dim=F;
|
|
P->ffn_x_k=rnd((size_t)C,0.5f);
|
|
P->ffn_key_w=rnd((size_t)C*F,0.3f);
|
|
P->ffn_value_w=rnd((size_t)F*C,0.3f);
|
|
}
|
|
static void mkG(rx_cm_grads*G,int C,int F){
|
|
G->ffn_x_k=zeros((size_t)C); G->ffn_key_w=zeros((size_t)C*F); G->ffn_value_w=zeros((size_t)F*C);
|
|
}
|
|
|
|
static void run_case(const char*name,int B,int T,int C,int F,uint64_t seed){
|
|
sm_state=seed; size_t BT=(size_t)B*T, tot=BT*(size_t)C;
|
|
rx_cm_params P; memset(&P,0,sizeof(P)); mkP(&P,C,F);
|
|
float*x=rnd(tot,1.0f),*xprev=rnd((size_t)B*C,1.0f),*dy=rnd(tot,1.0f);
|
|
float*out_c=(float*)malloc(tot*4),*out_g=(float*)malloc(tot*4);
|
|
rx_channelmix_forward_seq(&P,x,xprev,out_c,B,T);
|
|
if(channelmix_forward_cuda(&P,x,xprev,out_g,B,T)){printf(" %s GPU fwd ERR\n",name);g_fail=1;return;}
|
|
rx_cm_grads Gc,Gg; mkG(&Gc,C,F); mkG(&Gg,C,F);
|
|
float*dx_c=zeros(tot),*dx_g=zeros(tot),*dxp_c=zeros((size_t)B*C),*dxp_g=zeros((size_t)B*C);
|
|
rx_channelmix_backward(&P,x,xprev,dy,&Gc,dx_c,dxp_c,B,T);
|
|
if(channelmix_backward_cuda(&P,x,xprev,dy,&Gg,dx_g,dxp_g,B,T)){printf(" %s GPU bwd ERR\n",name);g_fail=1;return;}
|
|
printf(" %s (B=%d T=%d C=%d F=%d)\n",name,B,T,C,F);
|
|
chk("out",out_g,out_c,tot); chk("dx",dx_g,dx_c,tot); chk("d_xprev",dxp_g,dxp_c,(size_t)B*C);
|
|
chk("ffn_x_k",Gg.ffn_x_k,Gc.ffn_x_k,(size_t)C);
|
|
chk("ffn_key_w",Gg.ffn_key_w,Gc.ffn_key_w,(size_t)C*F);
|
|
chk("ffn_value_w",Gg.ffn_value_w,Gc.ffn_value_w,(size_t)F*C);
|
|
}
|
|
|
|
static void bench(int B,int T,int C,int F,int iters){
|
|
sm_state=99; size_t BT=(size_t)B*T, tot=BT*(size_t)C;
|
|
rx_cm_params P; memset(&P,0,sizeof(P)); mkP(&P,C,F);
|
|
float*x=rnd(tot,1.0f),*xprev=rnd((size_t)B*C,1.0f),*out=(float*)malloc(tot*4);
|
|
channelmix_forward_cuda(&P,x,xprev,out,B,T); /* warmup */
|
|
struct timespec t0,t1; clock_gettime(CLOCK_MONOTONIC,&t0);
|
|
for(int i=0;i<iters;i++) channelmix_forward_cuda(&P,x,xprev,out,B,T);
|
|
clock_gettime(CLOCK_MONOTONIC,&t1);
|
|
double sec=(t1.tv_sec-t0.tv_sec)+(t1.tv_nsec-t0.tv_nsec)*1e-9, toks=(double)BT*iters;
|
|
printf(" channelmix fwd B=%d T=%d C=%d F=%d: %d it in %.4fs -> %.0f tokens/sec (one K80 die; incl H2D+alloc)\n",
|
|
B,T,C,F,iters,sec,toks/sec);
|
|
}
|
|
|
|
int main(void){
|
|
int dev=0; cudaSetDevice(dev); cudaDeviceProp p; cudaGetDeviceProperties(&p,dev);
|
|
printf("=== GPU channel-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", 1,5,8,32, 11);
|
|
run_case("mid", 2,16,128,512,37);
|
|
run_case("C768", 4,32,768,3072,41);
|
|
printf("\n=== throughput (channel-mix forward, one die) ===\n");
|
|
bench(8,512,768,3072,20);
|
|
printf("\n%s\n", g_fail?"RESULT: FAIL":"RESULT: PASS");
|
|
return g_fail;
|
|
}
|