86 lines
4.3 KiB
Plaintext
86 lines
4.3 KiB
Plaintext
/* SPDX-License-Identifier: AGPL-3.0-or-later */
|
|
/* Copyright (C) 2026 David (Corey) Herrera */
|
|
/*
|
|
* test_layernorm_cuda.cu — parity harness for GPU LayerNorm (cuda_kernels/
|
|
* layernorm.cu) vs the CPU oracle rx_layernorm_forward_seq /
|
|
* rx_layernorm_backward (rwkvx/rx_pretrain.c), on GPU 0. Fixed-seed fixtures
|
|
* incl. the real config C=768. Parity metric: max-abs-error / tensor-scale
|
|
* (inf-norm), bar < 1e-3 (immune to signed-reduction near-zero entries).
|
|
* Covers out, dx, dw, db.
|
|
*/
|
|
#include <cuda_runtime.h>
|
|
#include <stdio.h>
|
|
#include <stdlib.h>
|
|
#include <stdint.h>
|
|
#include <string.h>
|
|
#include <math.h>
|
|
#include <time.h>
|
|
|
|
#include "cuda_kernels/layernorm.cu"
|
|
|
|
extern "C" int rx_layernorm_forward_seq(const float*,const float*,const float*,float*,int,int,int);
|
|
extern "C" int rx_layernorm_backward(const float*,const float*,const float*,const float*,
|
|
float*,float*,float*,int,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(" %-8s rel=%.3e maxabs=%.3e scale=%.3e [%s]\n",tname,rel,maxabs,scale,st);
|
|
}
|
|
|
|
static void run_case(const char*name,int B,int T,int C,uint64_t seed){
|
|
sm_state=seed; size_t BT=(size_t)B*T, tot=BT*(size_t)C;
|
|
float*x=rnd(tot,1.0f),*w=rndln((size_t)C),*b=rnd((size_t)C,0.2f),*dy=rnd(tot,1.0f);
|
|
float*out_c=(float*)malloc(tot*4),*out_g=(float*)malloc(tot*4);
|
|
rx_layernorm_forward_seq(x,w,b,out_c,B,T,C);
|
|
if(layernorm_forward_cuda(x,w,b,out_g,B,T,C)){printf(" %s GPU fwd ERR\n",name);g_fail=1;return;}
|
|
float*dx_c=zeros(tot),*dx_g=zeros(tot);
|
|
float*dw_c=zeros((size_t)C),*dw_g=zeros((size_t)C),*db_c=zeros((size_t)C),*db_g=zeros((size_t)C);
|
|
rx_layernorm_backward(x,w,b,dy,dx_c,dw_c,db_c,B,T,C);
|
|
if(layernorm_backward_cuda(x,w,dy,dx_g,dw_g,db_g,B,T,C)){printf(" %s GPU bwd ERR\n",name);g_fail=1;return;}
|
|
printf(" %s (B=%d T=%d C=%d)\n",name,B,T,C);
|
|
chk("out",out_g,out_c,tot); chk("dx",dx_g,dx_c,tot);
|
|
chk("dw",dw_g,dw_c,(size_t)C); chk("db",db_g,db_c,(size_t)C);
|
|
free(x);free(w);free(b);free(dy);free(out_c);free(out_g);
|
|
free(dx_c);free(dx_g);free(dw_c);free(dw_g);free(db_c);free(db_g);
|
|
}
|
|
|
|
static void bench(int B,int T,int C,int iters){
|
|
sm_state=99; size_t BT=(size_t)B*T, tot=BT*(size_t)C;
|
|
float*x=rnd(tot,1.0f),*w=rndln((size_t)C),*b=rnd((size_t)C,0.2f),*out=(float*)malloc(tot*4);
|
|
layernorm_forward_cuda(x,w,b,out,B,T,C); /* warmup */
|
|
struct timespec t0,t1; clock_gettime(CLOCK_MONOTONIC,&t0);
|
|
for(int i=0;i<iters;i++) layernorm_forward_cuda(x,w,b,out,B,T,C);
|
|
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(" layernorm fwd B=%d T=%d C=%d: %d it in %.4fs -> %.0f tokens/sec (one K80 die; incl H2D+alloc)\n",
|
|
B,T,C,iters,sec,toks/sec);
|
|
free(x);free(w);free(b);free(out);
|
|
}
|
|
|
|
int main(void){
|
|
int dev=0; cudaSetDevice(dev); cudaDeviceProp p; cudaGetDeviceProperties(&p,dev);
|
|
printf("=== GPU LayerNorm 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, 11);
|
|
run_case("mid", 2,16,128,37);
|
|
run_case("C768", 4,64,768,41);
|
|
printf("\n=== throughput (LayerNorm forward, one die) ===\n");
|
|
bench(8,512,768,50);
|
|
printf("\n%s\n", g_fail?"RESULT: FAIL":"RESULT: PASS");
|
|
return g_fail;
|
|
}
|