166 lines
8.5 KiB
Plaintext
166 lines
8.5 KiB
Plaintext
/* SPDX-License-Identifier: AGPL-3.0-or-later */
|
|
/* Copyright (C) 2026 David (Corey) Herrera */
|
|
/*
|
|
* channelmix.cu — RWKV-7 channel-mix (squared-ReLU FFN) fwd+bwd on the K80 (sm_37).
|
|
*
|
|
* GPU counterpart of the CPU oracle rx_channelmix_forward_seq /
|
|
* rx_channelmix_backward (rwkvx/rx_pretrain.c). EXCLUDES the pre-block ln2 and
|
|
* the residual add (those live in the block). Per token:
|
|
* sh = xx + (prev - xx) * ffn_x_k (token shift / lerp)
|
|
* fk = sh @ ffn_key_w ([C] -> [F])
|
|
* k = relu(fk)^2 (SQUARED ReLU activation)
|
|
* out = k @ ffn_value_w ([F] -> [C])
|
|
* prev at t=0 is x_prev_init (or zero); otherwise the previous token's input.
|
|
*
|
|
* Both matmuls go through cuBLAS SGEMM over the whole [B*T,*] token batch via the
|
|
* row-major->column-major rm_gemm() helper (swap A/B, pass N,M,K; same convention
|
|
* as fc_cuda.c / timemix.cu). Token-shift + squared-ReLU are elementwise kernels.
|
|
* The token-shift backward folds each token's own (1-lerp) term and the NEXT
|
|
* token's lerp term in one pass (no dx race); param grads use atomicAdd; weight
|
|
* grads fall out of single batched GEMMs (beta=0 == full token-sum).
|
|
*/
|
|
#include <cuda_runtime.h>
|
|
#include <cublas_v2.h>
|
|
#include <stdio.h>
|
|
#include <stdlib.h>
|
|
#include <string.h>
|
|
#include <math.h>
|
|
|
|
#include "rwkvx/rx_pretrain.h"
|
|
|
|
#define TCK(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 TBK(call) do{ cublasStatus_t s_=(call); if(s_!=CUBLAS_STATUS_SUCCESS){ \
|
|
fprintf(stderr,"cuBLAS %s:%d status %d\n",__FILE__,__LINE__,(int)s_); return -1; } }while(0)
|
|
|
|
#define TPB 256
|
|
static inline int nblk(size_t n){ return (int)((n+TPB-1)/TPB); }
|
|
|
|
/* ---- row-major GEMM via cuBLAS: C[M,N] = op(A) @ op(B) ---- */
|
|
static cublasStatus_t rm_gemm(cublasHandle_t H,int tA,int tB,int M,int N,int K,
|
|
float alpha,const float*A,const float*B,float beta,float*C){
|
|
int lda = tA? M : K;
|
|
int ldb = tB? K : N;
|
|
return cublasSgemm(H, tB?CUBLAS_OP_T:CUBLAS_OP_N, tA?CUBLAS_OP_T:CUBLAS_OP_N,
|
|
N, M, K, &alpha, B, ldb, A, lda, &beta, C, N);
|
|
}
|
|
|
|
/* ---- elementwise kernels ---- */
|
|
__global__ void k_cm_shift(const float*x,const float*xprev,const float*x_k,
|
|
float*sh,int B,int T,int C){
|
|
size_t g=(size_t)blockIdx.x*blockDim.x+threadIdx.x; size_t tot=(size_t)B*T*C; if(g>=tot)return;
|
|
int i=g%C; size_t token=g/C; int t=(int)(token%T); size_t b=token/T;
|
|
float xx=x[g];
|
|
float prev=(t==0)?(xprev?xprev[b*C+i]:0.0f):x[g-C];
|
|
sh[g]=xx+(prev-xx)*x_k[i];
|
|
}
|
|
__global__ void k_sqrelu(const float*fk,float*ka,size_t n){
|
|
size_t g=(size_t)blockIdx.x*blockDim.x+threadIdx.x; if(g>=n)return;
|
|
float v=fk[g]>0.0f?fk[g]:0.0f; ka[g]=v*v;
|
|
}
|
|
/* squared-relu backward: k=relu(fk)^2 => dfk = dk * 2*relu(fk) */
|
|
__global__ void k_sqrelu_bwd(const float*fk,const float*dka,float*dfk,size_t n){
|
|
size_t g=(size_t)blockIdx.x*blockDim.x+threadIdx.x; if(g>=n)return;
|
|
float r=fk[g]>0.0f?fk[g]:0.0f; dfk[g]=dka[g]*2.0f*r;
|
|
}
|
|
/* token-shift backward: sh = xx + (prev-xx)*x_k.
|
|
* dx[t] = dsh[t]*(1-x_k) + dsh[t+1]*x_k ; d_xprev_init = dsh[t=0]*x_k. */
|
|
__global__ void k_cm_shift_bwd(const float*x,const float*xprev,const float*dsh,
|
|
const float*x_k,float*G_x_k,float*dx,float*dxprev,int B,int T,int C){
|
|
size_t g=(size_t)blockIdx.x*blockDim.x+threadIdx.x; size_t tot=(size_t)B*T*C; if(g>=tot)return;
|
|
int i=g%C; size_t token=g/C; int t=(int)(token%T); size_t b=token/T;
|
|
float xx=x[g];
|
|
float prev=(t==0)?(xprev?xprev[b*C+i]:0.0f):x[g-C];
|
|
atomicAdd(&G_x_k[i], dsh[g]*(prev-xx));
|
|
float own=dsh[g]*(1.0f-x_k[i]);
|
|
float nxt=0.0f;
|
|
if(t<T-1) nxt=dsh[g+C]*x_k[i];
|
|
dx[g]=own+nxt;
|
|
if(t==0 && dxprev) dxprev[b*C+i]=dsh[g]*x_k[i];
|
|
}
|
|
|
|
/* ---- upload helper ---- */
|
|
static float* cup(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 channelmix_forward_cuda(const rx_cm_params*P,const float*x,
|
|
const float*x_prev_init,float*out,int B,int T){
|
|
int C=P->c, F=P->ffn_dim; size_t BT=(size_t)B*T, tot=BT*(size_t)C, totF=BT*(size_t)F;
|
|
cublasHandle_t H; if(cublasCreate(&H)!=CUBLAS_STATUS_SUCCESS) return -1;
|
|
float*dx=cup(x,tot),*dxk=cup(P->ffn_x_k,(size_t)C);
|
|
float*dkw=cup(P->ffn_key_w,(size_t)C*F),*dvw=cup(P->ffn_value_w,(size_t)F*C);
|
|
float*dxprev = x_prev_init? cup(x_prev_init,(size_t)B*C) : NULL;
|
|
float*sh=NULL,*fk=NULL,*ka=NULL,*dout=NULL;
|
|
TCK(cudaMalloc(&sh,tot*sizeof(float))); TCK(cudaMalloc(&fk,totF*sizeof(float)));
|
|
TCK(cudaMalloc(&ka,totF*sizeof(float))); TCK(cudaMalloc(&dout,tot*sizeof(float)));
|
|
k_cm_shift<<<nblk(tot),TPB>>>(dx,dxprev,dxk,sh,B,T,C);
|
|
TBK(rm_gemm(H,0,0,BT,F,C,1.0f,sh,dkw,0.0f,fk));
|
|
k_sqrelu<<<nblk(totF),TPB>>>(fk,ka,totF);
|
|
TBK(rm_gemm(H,0,0,BT,C,F,1.0f,ka,dvw,0.0f,dout));
|
|
TCK(cudaDeviceSynchronize());
|
|
TCK(cudaMemcpy(out,dout,tot*sizeof(float),cudaMemcpyDeviceToHost));
|
|
cudaFree(dx);cudaFree(dxk);cudaFree(dkw);cudaFree(dvw);if(dxprev)cudaFree(dxprev);
|
|
cudaFree(sh);cudaFree(fk);cudaFree(ka);cudaFree(dout);
|
|
cublasDestroy(H);
|
|
return 0;
|
|
}
|
|
|
|
/* ---- host-facing backward parity entry ---- */
|
|
extern "C" int channelmix_backward_cuda(const rx_cm_params*P,const float*x,
|
|
const float*x_prev_init,const float*dy,rx_cm_grads*G,float*dx,float*d_xprev_init,
|
|
int B,int T){
|
|
int C=P->c, F=P->ffn_dim; size_t BT=(size_t)B*T, tot=BT*(size_t)C, totF=BT*(size_t)F;
|
|
cublasHandle_t H; if(cublasCreate(&H)!=CUBLAS_STATUS_SUCCESS) return -1;
|
|
float*dxin=cup(x,tot),*dxk=cup(P->ffn_x_k,(size_t)C);
|
|
float*dkw=cup(P->ffn_key_w,(size_t)C*F),*dvw=cup(P->ffn_value_w,(size_t)F*C);
|
|
float*ddy=cup(dy,tot);
|
|
float*dxprev = x_prev_init? cup(x_prev_init,(size_t)B*C) : NULL;
|
|
|
|
/* device grads (zeroed) */
|
|
float*G_xk=NULL,*G_kw=NULL,*G_vw=NULL;
|
|
TCK(cudaMalloc(&G_xk,(size_t)C*sizeof(float))); TCK(cudaMemset(G_xk,0,(size_t)C*sizeof(float)));
|
|
TCK(cudaMalloc(&G_kw,(size_t)C*F*sizeof(float))); TCK(cudaMemset(G_kw,0,(size_t)C*F*sizeof(float)));
|
|
TCK(cudaMalloc(&G_vw,(size_t)F*C*sizeof(float))); TCK(cudaMemset(G_vw,0,(size_t)F*C*sizeof(float)));
|
|
|
|
/* scratch + recompute forward (sh, fk, ka) */
|
|
float*sh=NULL,*fk=NULL,*ka=NULL,*dka=NULL,*dfk=NULL,*dsh=NULL,*ddx=NULL,*ddxprev=NULL;
|
|
TCK(cudaMalloc(&sh,tot*sizeof(float))); TCK(cudaMalloc(&fk,totF*sizeof(float)));
|
|
TCK(cudaMalloc(&ka,totF*sizeof(float))); TCK(cudaMalloc(&dka,totF*sizeof(float)));
|
|
TCK(cudaMalloc(&dfk,totF*sizeof(float))); TCK(cudaMalloc(&dsh,tot*sizeof(float)));
|
|
TCK(cudaMalloc(&ddx,tot*sizeof(float)));
|
|
if(d_xprev_init){ TCK(cudaMalloc(&ddxprev,(size_t)B*C*sizeof(float)));
|
|
TCK(cudaMemset(ddxprev,0,(size_t)B*C*sizeof(float))); }
|
|
|
|
k_cm_shift<<<nblk(tot),TPB>>>(dxin,dxprev,dxk,sh,B,T,C);
|
|
TBK(rm_gemm(H,0,0,BT,F,C,1.0f,sh,dkw,0.0f,fk));
|
|
k_sqrelu<<<nblk(totF),TPB>>>(fk,ka,totF);
|
|
|
|
/* value proj backward: out = ka @ Wv */
|
|
TBK(rm_gemm(H,0,1,BT,F,C,1.0f,ddy,dvw,0.0f,dka)); /* dka = dy @ Wv^T */
|
|
TBK(rm_gemm(H,1,0,F,C,BT,1.0f,ka,ddy,0.0f,G_vw)); /* G_Wv = ka^T @ dy */
|
|
/* squared-relu backward */
|
|
k_sqrelu_bwd<<<nblk(totF),TPB>>>(fk,dka,dfk,totF);
|
|
/* key proj backward: fk = sh @ Wk */
|
|
TBK(rm_gemm(H,0,1,BT,C,F,1.0f,dfk,dkw,0.0f,dsh)); /* dsh = dfk @ Wk^T */
|
|
TBK(rm_gemm(H,1,0,C,F,BT,1.0f,sh,dfk,0.0f,G_kw)); /* G_Wk = sh^T @ dfk */
|
|
/* token-shift backward -> dx, d_xprev_init, G_x_k */
|
|
k_cm_shift_bwd<<<nblk(tot),TPB>>>(dxin,dxprev,dsh,dxk,G_xk,ddx,ddxprev,B,T,C);
|
|
TCK(cudaDeviceSynchronize());
|
|
|
|
/* download */
|
|
TCK(cudaMemcpy(dx,ddx,tot*sizeof(float),cudaMemcpyDeviceToHost));
|
|
if(d_xprev_init) TCK(cudaMemcpy(d_xprev_init,ddxprev,(size_t)B*C*sizeof(float),cudaMemcpyDeviceToHost));
|
|
if(G->ffn_x_k) TCK(cudaMemcpy(G->ffn_x_k,G_xk,(size_t)C*sizeof(float),cudaMemcpyDeviceToHost));
|
|
if(G->ffn_key_w) TCK(cudaMemcpy(G->ffn_key_w,G_kw,(size_t)C*F*sizeof(float),cudaMemcpyDeviceToHost));
|
|
if(G->ffn_value_w)TCK(cudaMemcpy(G->ffn_value_w,G_vw,(size_t)F*C*sizeof(float),cudaMemcpyDeviceToHost));
|
|
|
|
cudaFree(dxin);cudaFree(dxk);cudaFree(dkw);cudaFree(dvw);cudaFree(ddy);if(dxprev)cudaFree(dxprev);
|
|
cudaFree(G_xk);cudaFree(G_kw);cudaFree(G_vw);
|
|
cudaFree(sh);cudaFree(fk);cudaFree(ka);cudaFree(dka);cudaFree(dfk);cudaFree(dsh);
|
|
cudaFree(ddx);if(ddxprev)cudaFree(ddxprev);
|
|
cublasDestroy(H);
|
|
return 0;
|
|
}
|