100 lines
5.8 KiB
Plaintext
100 lines
5.8 KiB
Plaintext
/* SPDX-License-Identifier: AGPL-3.0-or-later */
|
|
/* Copyright (C) 2026 David (Corey) Herrera */
|
|
/*
|
|
* test_embed_ce_cuda.cu — isolated parity for the three new leaf kernels of the
|
|
* full-model step, vs the CPU oracle (rwkvx/rx_pretrain.c), on GPU 0:
|
|
* - embedding gather fwd + scatter-add bwd (rx_embedding_forward/backward)
|
|
* - LM head fwd + bwd (cuBLAS) (rx_head_forward/backward)
|
|
* - fused log-softmax cross-entropy fwd/bwd (rx_cross_entropy_*)
|
|
* Reuses the device kernels from cuda_kernels/fullmodel.cu. inf-norm bar 1e-3.
|
|
* Small vocab so tokens REPEAT — exercises the scatter-add row collisions.
|
|
*/
|
|
#include <cuda_runtime.h>
|
|
#include <cublas_v2.h>
|
|
#include <stdio.h>
|
|
#include <stdlib.h>
|
|
#include <stdint.h>
|
|
#include <string.h>
|
|
#include <math.h>
|
|
|
|
#include "cuda_kernels/fullmodel.cu"
|
|
|
|
static uint64_t S; static uint64_t nx(void){ uint64_t z=(S+=0x9E3779B97F4A7C15ULL);
|
|
z=(z^(z>>30))*0xBF58476D1CE4E5B9ULL; z=(z^(z>>27))*0x94D049BB133111EBULL; return z^(z>>31); }
|
|
static float uf(void){ return (float)((nx()>>11)*(1.0/9007199254740992.0)); }
|
|
static float sy(float s){ return (2.0f*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]=sy(s); return p; }
|
|
static int g_fail=0;
|
|
static void chk(const char*t,const float*g,const float*c,size_t n,double bar){
|
|
double ma=0,sc=0; for(size_t i=0;i<n;i++){ double d=fabs((double)g[i]-c[i]); if(d>ma)ma=d;
|
|
double a=fabs((double)c[i]); if(a>sc)sc=a; }
|
|
double rel=ma/(sc+1e-30); if(rel>=bar)g_fail=1;
|
|
printf(" %-10s rel=%.3e maxabs=%.3e scale=%.3e [%s]\n",t,rel,ma,sc,rel<bar?"PASS":"FAIL"); }
|
|
|
|
int main(void){
|
|
cudaSetDevice(0); cublasHandle_t H; cublasCreate(&H);
|
|
int T=24,C=48,V=16; /* V<T so tokens repeat -> scatter-add collisions */
|
|
size_t TC=(size_t)T*C, TV=(size_t)T*V, VC=(size_t)V*C, CV=(size_t)C*V;
|
|
printf("=== embedding / head / cross-entropy isolated parity (GPU 0) ===\n");
|
|
S=7;
|
|
float*emb=rnd(VC,0.5f),*head=rnd(CV,0.3f),*x=rnd(TC,1.0f);
|
|
int*tok=(int*)malloc(T*4),*tgt=(int*)malloc(T*4);
|
|
for(int t=0;t<T;t++){ tok[t]=(int)(nx()%V); tgt[t]=(int)(nx()%V); }
|
|
|
|
/* ---- embedding fwd ---- */
|
|
float*eo_c=(float*)malloc(TC*4); rx_embedding_forward(emb,tok,T,C,eo_c);
|
|
float *demb,*eo_g; int *dtok; cudaMalloc(&demb,VC*4);cudaMalloc(&dtok,T*4);cudaMalloc(&eo_g,TC*4);
|
|
cudaMemcpy(demb,emb,VC*4,cudaMemcpyHostToDevice); cudaMemcpy(dtok,tok,T*4,cudaMemcpyHostToDevice);
|
|
fm_k_emb_fwd<<<fm_nblk(TC),256>>>(demb,dtok,eo_g,T,C); cudaDeviceSynchronize();
|
|
float*tmp=(float*)malloc(TC*4); cudaMemcpy(tmp,eo_g,TC*4,cudaMemcpyDeviceToHost);
|
|
printf(" [embedding]\n"); chk("emb_fwd",tmp,eo_c,TC,1e-3);
|
|
/* ---- embedding bwd (scatter-add) ---- */
|
|
float*dout=rnd(TC,1.0f); float*demb_c=(float*)malloc(VC*4);
|
|
rx_embedding_backward(tok,T,C,V,dout,demb_c);
|
|
float*ddout,*dgemb; cudaMalloc(&ddout,TC*4);cudaMalloc(&dgemb,VC*4);
|
|
cudaMemcpy(ddout,dout,TC*4,cudaMemcpyHostToDevice); cudaMemset(dgemb,0,VC*4);
|
|
fm_k_emb_bwd<<<fm_nblk(TC),256>>>(dtok,ddout,dgemb,T,C); cudaDeviceSynchronize();
|
|
float*dembg=(float*)malloc(VC*4); cudaMemcpy(dembg,dgemb,VC*4,cudaMemcpyDeviceToHost);
|
|
chk("emb_bwd",dembg,demb_c,VC,1e-3);
|
|
|
|
/* ---- head fwd/bwd ---- */
|
|
float*lg_c=(float*)malloc(TV*4); rx_head_forward(x,head,T,C,V,lg_c);
|
|
float*dx_c=(float*)malloc(TC*4),*dh_c=(float*)malloc(CV*4);
|
|
float*dl=rnd(TV,1.0f); rx_head_backward(x,head,dl,T,C,V,dx_c,dh_c);
|
|
float *dxg,*dhead,*dxdev,*dlg,*dlogit;
|
|
cudaMalloc(&dxg,TC*4);cudaMalloc(&dhead,CV*4);cudaMalloc(&dxdev,TC*4);cudaMalloc(&dlg,TV*4);cudaMalloc(&dlogit,TV*4);
|
|
cudaMemcpy(dxdev,x,TC*4,cudaMemcpyHostToDevice); cudaMemcpy(dhead,head,CV*4,cudaMemcpyHostToDevice);
|
|
cudaMemcpy(dlg,dl,TV*4,cudaMemcpyHostToDevice);
|
|
rm_gemm_tm(H,0,0,T,V,C,1.0f,dxdev,dhead,0.0f,dlogit); /* logits */
|
|
float*lg_g=(float*)malloc(TV*4); cudaMemcpy(lg_g,dlogit,TV*4,cudaMemcpyDeviceToHost);
|
|
printf(" [head]\n"); chk("head_fwd",lg_g,lg_c,TV,1e-3);
|
|
float*dheadg2; cudaMalloc(&dheadg2,CV*4);
|
|
rm_gemm_tm(H,1,0,C,V,T,1.0f,dxdev,dlg,0.0f,dheadg2); /* d_head=x^T@dl */
|
|
rm_gemm_tm(H,0,1,T,C,V,1.0f,dlg,dhead,0.0f,dxg); /* dx=dl@head^T */
|
|
float*dxh=(float*)malloc(TC*4),*dhh=(float*)malloc(CV*4);
|
|
cudaMemcpy(dxh,dxg,TC*4,cudaMemcpyDeviceToHost); cudaMemcpy(dhh,dheadg2,CV*4,cudaMemcpyDeviceToHost);
|
|
chk("head_dx",dxh,dx_c,TC,1e-3); chk("head_dW",dhh,dh_c,CV,1e-3);
|
|
|
|
/* ---- cross-entropy fwd/bwd (ignore=-1 and ignore case) ---- */
|
|
for(int trial=0;trial<2;trial++){
|
|
int ign = trial? tgt[3] : -1;
|
|
float*dlc=(float*)malloc(TV*4);
|
|
double loss_c=rx_cross_entropy_backward(lg_c,tgt,T,V,ign,dlc);
|
|
int nv=0; for(int t=0;t<T;t++) if(!(ign>=0&&tgt[t]==ign)) nv++;
|
|
int*dtgt; cudaMalloc(&dtgt,T*4); cudaMemcpy(dtgt,tgt,T*4,cudaMemcpyHostToDevice);
|
|
float*dlg2; cudaMalloc(&dlg2,TV*4); double*dlt; cudaMalloc(&dlt,(size_t)T*8);
|
|
float*dlogit_c; cudaMalloc(&dlogit_c,TV*4); cudaMemcpy(dlogit_c,lg_c,TV*4,cudaMemcpyHostToDevice);
|
|
fm_k_ce_bwd<<<fm_nblk(T),256>>>(dlogit_c,dtgt,dlg2,dlt,T,V,ign,nv); cudaDeviceSynchronize();
|
|
double*lth=(double*)malloc((size_t)T*8); cudaMemcpy(lth,dlt,(size_t)T*8,cudaMemcpyDeviceToHost);
|
|
double loss_g=0; for(int t=0;t<T;t++) loss_g+=lth[t]; loss_g=nv>0?loss_g/nv:0;
|
|
float*dlgh=(float*)malloc(TV*4); cudaMemcpy(dlgh,dlg2,TV*4,cudaMemcpyDeviceToHost);
|
|
printf(" [cross-entropy ignore=%d nv=%d] loss cpu=%.8f gpu=%.8f rel=%.3e\n",
|
|
ign,nv,loss_c,loss_g,fabs(loss_g-loss_c)/(fabs(loss_c)+1e-30));
|
|
if(fabs(loss_g-loss_c)/(fabs(loss_c)+1e-30)>=1e-3) g_fail=1;
|
|
chk("ce_dlogit",dlgh,dlc,TV,1e-3);
|
|
free(dlc);free(lth);free(dlgh);
|
|
}
|
|
printf("\n%s\n", g_fail?"RESULT: FAIL":"RESULT: PASS");
|
|
return g_fail;
|
|
}
|