113 lines
5.3 KiB
Plaintext
113 lines
5.3 KiB
Plaintext
/* SPDX-License-Identifier: AGPL-3.0-or-later */
|
|
/* Copyright (C) 2026 David (Corey) Herrera */
|
|
/*
|
|
* layernorm.cu — RWKV-x LayerNorm forward+backward on the K80 (sm_37).
|
|
*
|
|
* GPU counterpart of the CPU oracle rx_layernorm_forward_seq /
|
|
* rx_layernorm_backward (rwkvx/rx_pretrain.c). LayerNorm over the full channel
|
|
* dim C for each of the B*T tokens, eps=1e-5:
|
|
* mean = (1/C) sum_i x_i ; var = (1/C) sum_i (x_i-mean)^2
|
|
* rstd = 1/sqrt(var+eps) ; out_i = (x_i-mean)*rstd*w_i + b_i
|
|
*
|
|
* Numerics: the mean/variance AND both backward reductions (mean(dnorm),
|
|
* mean(dnorm*norm)) are cancellation-prone. One thread owns one token's entire
|
|
* C-length reduction, accumulating in DOUBLE — matching the CPU oracle's double
|
|
* discipline exactly. The backward RE-derives mean/var/rstd in double from x
|
|
* (the oracle does the same; it does not reuse a truncated fp32 cache), so the
|
|
* two paths are bit-for-bit in the reduction that matters. dw/db are summed
|
|
* across all tokens with atomicAdd (order-independent, well under the 1e-3 bar).
|
|
*
|
|
* Token layout [B][T][C] row-major == BT rows of length C (no reshape).
|
|
*/
|
|
#include <cuda_runtime.h>
|
|
#include <stdio.h>
|
|
#include <stdlib.h>
|
|
#include <math.h>
|
|
|
|
#include "rwkvx/rx_pretrain.h"
|
|
|
|
#define LCK(call) do{ cudaError_t e_=(call); if(e_!=cudaSuccess){ \
|
|
fprintf(stderr,"CUDA %s:%d %s\n",__FILE__,__LINE__,cudaGetErrorString(e_)); return -1; } }while(0)
|
|
|
|
#define RX_LN_EPS 1e-5
|
|
#define TPB 128
|
|
static inline int nblk(size_t n){ return (int)((n+TPB-1)/TPB); }
|
|
|
|
/* forward: one thread per token, double-accumulated mean/var.
|
|
* mean_o/rstd_o are optional fp32 stat caches (pass NULL to skip). */
|
|
__global__ void k_ln_fwd(const float*x,const float*w,const float*b,float*out,
|
|
float*mean_o,float*rstd_o,int C,size_t BT){
|
|
size_t tk=(size_t)blockIdx.x*blockDim.x+threadIdx.x; if(tk>=BT)return;
|
|
const float*xx=x+tk*C; float*o=out+tk*C;
|
|
double mean=0.0; for(int i=0;i<C;i++) mean+=xx[i]; mean/=C;
|
|
double var=0.0; for(int i=0;i<C;i++){ double d=(double)xx[i]-mean; var+=d*d; } var/=C;
|
|
double rstd=1.0/sqrt(var+RX_LN_EPS);
|
|
if(mean_o) mean_o[tk]=(float)mean;
|
|
if(rstd_o) rstd_o[tk]=(float)rstd;
|
|
for(int i=0;i<C;i++) o[i]=(float)(((double)xx[i]-mean)*rstd)*w[i]+b[i];
|
|
}
|
|
|
|
/* backward: one thread per token. Re-derives mean/var/rstd in double (matches
|
|
* oracle). dnorm=dy*w; mdn=mean(dnorm); mdnn=mean(dnorm*norm);
|
|
* dx = rstd*(dnorm - mdn - norm*mdnn). dw+=dy*norm; db+=dy (atomic). */
|
|
__global__ void k_ln_bwd(const float*x,const float*w,const float*dy,
|
|
float*dx,float*dw,float*db,int C,size_t BT){
|
|
size_t tk=(size_t)blockIdx.x*blockDim.x+threadIdx.x; if(tk>=BT)return;
|
|
const float*xx=x+tk*C; const float*dyt=dy+tk*C; float*dxt=dx+tk*C;
|
|
double mean=0.0; for(int i=0;i<C;i++) mean+=xx[i]; mean/=C;
|
|
double var=0.0; for(int i=0;i<C;i++){ double d=(double)xx[i]-mean; var+=d*d; } var/=C;
|
|
double rstd=1.0/sqrt(var+RX_LN_EPS);
|
|
double mdn=0.0, mdnn=0.0;
|
|
for(int i=0;i<C;i++){
|
|
double nv=((double)xx[i]-mean)*rstd;
|
|
double dnorm=(double)dyt[i]*(double)w[i];
|
|
mdn+=dnorm; mdnn+=dnorm*nv;
|
|
if(dw) atomicAdd(&dw[i], dyt[i]*(float)nv);
|
|
if(db) atomicAdd(&db[i], dyt[i]);
|
|
}
|
|
mdn/=C; mdnn/=C;
|
|
for(int i=0;i<C;i++){
|
|
double nv=((double)xx[i]-mean)*rstd;
|
|
double dnorm=(double)dyt[i]*(double)w[i];
|
|
dxt[i]=(float)(rstd*(dnorm - mdn - nv*mdnn));
|
|
}
|
|
}
|
|
|
|
/* ---- host->device upload helper (NULL passthrough) ---- */
|
|
static float* lup(const float*h,size_t n){
|
|
if(!h) return NULL; float*d=NULL; if(cudaMalloc(&d,n*sizeof(float))!=cudaSuccess) return NULL;
|
|
if(cudaMemcpy(d,h,n*sizeof(float),cudaMemcpyHostToDevice)!=cudaSuccess){cudaFree(d);return NULL;} return d; }
|
|
|
|
/* ---- host-facing forward parity entry ---- */
|
|
extern "C" int layernorm_forward_cuda(const float*x,const float*w,const float*b,
|
|
float*out,int B,int T,int C){
|
|
size_t BT=(size_t)B*T, tot=BT*(size_t)C;
|
|
float*dx=lup(x,tot),*dw=lup(w,(size_t)C),*db=lup(b,(size_t)C),*dout=NULL;
|
|
if(!dx||!dw||!db) return -1;
|
|
LCK(cudaMalloc(&dout,tot*sizeof(float)));
|
|
k_ln_fwd<<<nblk(BT),TPB>>>(dx,dw,db,dout,NULL,NULL,C,BT);
|
|
LCK(cudaDeviceSynchronize());
|
|
LCK(cudaMemcpy(out,dout,tot*sizeof(float),cudaMemcpyDeviceToHost));
|
|
cudaFree(dx);cudaFree(dw);cudaFree(db);cudaFree(dout);
|
|
return 0;
|
|
}
|
|
|
|
/* ---- host-facing backward parity entry ---- */
|
|
extern "C" int layernorm_backward_cuda(const float*x,const float*w,const float*dy,
|
|
float*dx,float*dw,float*db,int B,int T,int C){
|
|
size_t BT=(size_t)B*T, tot=BT*(size_t)C;
|
|
float* dxin=lup(x,tot),*dwt=lup(w,(size_t)C),*ddy=lup(dy,tot);
|
|
float*ddx=NULL,*ddw=NULL,*ddb=NULL;
|
|
if(!dxin||!dwt||!ddy) return -1;
|
|
LCK(cudaMalloc(&ddx,tot*sizeof(float)));
|
|
LCK(cudaMalloc(&ddw,(size_t)C*sizeof(float))); LCK(cudaMemset(ddw,0,(size_t)C*sizeof(float)));
|
|
LCK(cudaMalloc(&ddb,(size_t)C*sizeof(float))); LCK(cudaMemset(ddb,0,(size_t)C*sizeof(float)));
|
|
k_ln_bwd<<<nblk(BT),TPB>>>(dxin,dwt,ddy,ddx, dw?ddw:NULL, db?ddb:NULL, C,BT);
|
|
LCK(cudaDeviceSynchronize());
|
|
LCK(cudaMemcpy(dx,ddx,tot*sizeof(float),cudaMemcpyDeviceToHost));
|
|
if(dw) LCK(cudaMemcpy(dw,ddw,(size_t)C*sizeof(float),cudaMemcpyDeviceToHost));
|
|
if(db) LCK(cudaMemcpy(db,ddb,(size_t)C*sizeof(float),cudaMemcpyDeviceToHost));
|
|
cudaFree(dxin);cudaFree(dwt);cudaFree(ddy);cudaFree(ddx);cudaFree(ddw);cudaFree(ddb);
|
|
return 0;
|
|
}
|