David
/
Foundry
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
Foundry/cuda_kernels/fullmodel.cu

635 lines
36 KiB
Plaintext
Raw Blame History

/* SPDX-License-Identifier: AGPL-3.0-or-later */
/* Copyright (C) 2026 David (Corey) Herrera */
/*
* fullmodel.cu — FULL RWKV-x training step on the K80 (sm_37), device-resident.
*
* The Phase-2 completion piece: embedding gather (fwd + scatter-add bwd), LM head
* (fwd + bwd, cuBLAS), fused log-softmax cross-entropy (fwd + bwd, double-
* accumulated LSE/normalizer), and the WHOLE-model assembly:
*
* x0 = emb[token]
* for each interleaved block (RWKV-7 time-mix OR MoBA sparse-attention):
* x = block(x) (ln1 -> core -> +x -> ln2 -> cm -> +x1)
* h = ln_out(x) ; logits = h @ head ; loss = mean_t CE(logits[t], target[t])
*
* GPU counterpart of the CPU oracle rx_pretrain_forward / rx_pretrain_backward
* (rwkvx/rx_pretrain.c), B=1 single sequence. Head is NOT weight-tied to the
* embedding (foundry stores rwkv.head.weight separately) — each gets its own grad.
*
* REUSE: this is ONE translation unit that includes every already-merged block
* kernel via a rename/undef shim, so all their device helpers (tm_fwd_device,
* k_ln_fwd/bwd, k_cm_*, moba_k_*, the rm_gemm cuBLAS helper, the WKV7 recurrence)
* are called directly. Weights upload ONCE into resident device structs; the
* backward recomputes each block's sub-forward from the cached residual stream,
* exactly as the CPU oracle does. No edits to the merged component files.
*
* Build (K80 host, GPU 0, standalone — NOT the fc_cuda_ptx.h embed pipeline):
* gcc-11 ... -c rwkvx/rx_pretrain.c -o rx_pretrain_fm.o
* nvcc -arch=sm_37 -ccbin gcc-11 -isystem cuda_kernels/nvcc_shim -I. \
* test_fullmodel_cuda.cu rx_pretrain_fm.o -lcublas -lm
*/
#include <cuda_runtime.h>
#include <cublas_v2.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include <stdint.h>
/* ============================ include shim ============================
* Pull all merged block kernels into this single TU. Only nblk / rm_gemm
* (static) and TPB / TCK / TBK (macros) collide across files; every __global__
* kernel name is already unique. Rename the two statics per file and undef the
* macros between includes. wkv7.cu is pulled in by timemix.cu (only). */
#define nblk nblk_tm
#define rm_gemm rm_gemm_tm
#include "cuda_kernels/timemix.cu"
#undef nblk
#undef rm_gemm
#undef TPB
#undef TCK
#undef TBK
#undef RWDECAY_COEFF
#define nblk nblk_cm
#define rm_gemm rm_gemm_cm
#include "cuda_kernels/channelmix.cu"
#undef nblk
#undef rm_gemm
#undef TPB
#undef TCK
#undef TBK
#define nblk nblk_ln
#include "cuda_kernels/layernorm.cu"
#undef nblk
#undef TPB
#define rm_gemm rm_gemm_moba
#include "cuda_kernels/moba_attn.cu"
#undef rm_gemm
#include "rwkvx/rx_pretrain.h" /* guarded; brings the oracle decls + structs */
/* ============================ fullmodel-local kernels ============================ */
#define FM_TPB 256
static inline int fm_nblk(size_t n){ return (int)((n+FM_TPB-1)/FM_TPB); }
#define FCK(call) do{ cudaError_t e_=(call); if(e_!=cudaSuccess){ \
fprintf(stderr,"CUDA %s:%d %s\n",__FILE__,__LINE__,cudaGetErrorString(e_)); exit(1);} }while(0)
#define FBK(call) do{ cublasStatus_t s_=(call); if(s_!=CUBLAS_STATUS_SUCCESS){ \
fprintf(stderr,"cuBLAS %s:%d status %d\n",__FILE__,__LINE__,(int)s_); exit(1);} }while(0)
__global__ void fm_k_emb_fwd(const float*emb,const int*tok,float*out,int T,int C){
size_t g=(size_t)blockIdx.x*blockDim.x+threadIdx.x; size_t tot=(size_t)T*C; if(g>=tot)return;
int i=(int)(g%C); int t=(int)(g/C); out[g]=emb[(size_t)tok[t]*C+i];
}
/* scatter-add: one thread per (t,i) atomicAdds into the token's embedding row.
* Row collisions (a token repeated in the sequence) are resolved by atomicAdd. */
__global__ void fm_k_emb_bwd(const int*tok,const float*dout,float*demb,int T,int C){
size_t g=(size_t)blockIdx.x*blockDim.x+threadIdx.x; size_t tot=(size_t)T*C; if(g>=tot)return;
int i=(int)(g%C); int t=(int)(g/C); atomicAdd(&demb[(size_t)tok[t]*C+i], dout[g]);
}
__global__ void fm_k_add(const float*a,const float*b,float*o,size_t n){
size_t g=(size_t)blockIdx.x*blockDim.x+threadIdx.x; if(g>=n)return; o[g]=a[g]+b[g]; }
__global__ void fm_k_addi(float*a,const float*b,size_t n){
size_t g=(size_t)blockIdx.x*blockDim.x+threadIdx.x; if(g>=n)return; a[g]+=b[g]; }
/* fused log-softmax cross-entropy, one thread per token (double LSE + normalizer).
* Writes dlogits=(softmax-onehot)/nv (ignored rows zeroed) and per-token loss
* (lse - logit[target]) into loss_t[T] (double, host reduces -> /nv). sm_37 has
* no double atomicAdd, so we avoid a global reduction entirely. */
__global__ void fm_k_ce_bwd(const float*logits,const int*targets,float*dlogits,
double*loss_t,int T,int V,int ignore_index,int nv){
int t=blockIdx.x*blockDim.x+threadIdx.x; if(t>=T)return;
int tg=targets[t]; const float*lt=logits+(size_t)t*V; float*dl=dlogits+(size_t)t*V;
if(ignore_index>=0 && tg==ignore_index){ for(int j=0;j<V;j++) dl[j]=0.0f; loss_t[t]=0.0; return; }
double mx=-1e300; for(int j=0;j<V;j++){ double d=(double)lt[j]; if(d>mx)mx=d; }
double se=0.0; for(int j=0;j<V;j++) se+=exp((double)lt[j]-mx);
double lse=mx+log(se);
double inv= nv>0 ? 1.0/(double)nv : 0.0;
for(int j=0;j<V;j++){ double sm=exp((double)lt[j]-lse); double g=sm-(j==tg?1.0:0.0); dl[j]=(float)(g*inv); }
loss_t[t]=lse-(double)lt[tg];
}
/* ============================ resident device model ============================ */
typedef struct { float *x_k,*key_w,*value_w; } fm_cmw; /* channel-mix weights/grads */
typedef struct {
float *ln1_w,*ln1_b,*ln2_w,*ln2_b;
float *Wq,*Wk,*Wv,*Wo;
fm_cmw cm;
rx_attn_params P; /* host param mirror (dims + a live pointer set) */
} fm_attn_blk;
typedef struct {
tm_dw D, Gd; /* time-mix weights + grads (device) */
fm_cmw cm, cmG; /* channel-mix weights + grads */
float *ln1_w,*ln1_b,*ln2_w,*ln2_b;
float *g_ln1_w,*g_ln1_b,*g_ln2_w,*g_ln2_b;
rx_tm_params tmP; /* host mirror for dims/layer_id */
rx_cm_params cmP;
} fm_rwkv_blk;
/* flat grad registry for name-keyed parity download */
#define FM_MAXG 4096
typedef struct {
int C,V,nH,N,F,Dw,Da,Dg,Dv,n_rwkv,n_attn,n_blocks,T;
int *btype,*bidx; /* host [n_blocks] */
float *emb,*head,*lnout_w,*lnout_b; /* device weights */
float *g_emb,*g_head,*g_lnout_w,*g_lnout_b; /* device grads */
fm_rwkv_blk *rw; fm_attn_blk *at;
/* activations (device) */
int *tok;
int *dtg; /* resident target ids (opt#3: stable ptr for graph replay) */
float *block_in; /* (n_blocks+1)*T*C */
float *v_first; /* T*C */
float *lnout; /* T*C */
float *logits; /* T*V */
double*loss_t; /* T */
cublasHandle_t H;
/* grad registry */
char gname[FM_MAXG][48];
float *gptr[FM_MAXG]; int gn[FM_MAXG]; int ng;
} fm_model;
static float* fm_up(const float*h,size_t n){
float*d=NULL; FCK(cudaMalloc(&d,n*sizeof(float)));
if(h) FCK(cudaMemcpy(d,h,n*sizeof(float),cudaMemcpyHostToDevice));
else FCK(cudaMemset(d,0,n*sizeof(float)));
return d;
}
static float* fm_upg(fm_model*m,const char*name,int n){ /* alloc a zeroed device grad + register */
float*d=NULL; FCK(cudaMalloc(&d,(size_t)n*sizeof(float))); FCK(cudaMemset(d,0,(size_t)n*sizeof(float)));
snprintf(m->gname[m->ng],48,"%s",name); m->gptr[m->ng]=d; m->gn[m->ng]=n; m->ng++;
return d;
}
static float* fm_gfind(fm_model*m,const char*name,int*n){
for(int i=0;i<m->ng;i++) if(!strcmp(m->gname[i],name)){ if(n)*n=m->gn[i]; return m->gptr[i]; }
return NULL;
}
/* upload a whole model from the CPU oracle's rx_ptmodel (same names/shapes). */
static fm_model* fm_build(rx_ptmodel *cm, int T){
rx_pt_config pc=rx_ptmodel_config(cm);
fm_model*m=(fm_model*)calloc(1,sizeof(fm_model));
m->C=pc.n_embd; m->V=pc.vocab_size; m->nH=pc.n_head; m->N=pc.head_size;
m->F=pc.ffn_dim; m->Dw=pc.Dw; m->Da=pc.Da; m->Dg=pc.Dg; m->Dv=pc.Dv;
m->n_rwkv=pc.n_rwkv_layer; m->n_attn=pc.n_attn_layer; m->n_blocks=m->n_rwkv+m->n_attn; m->T=T;
int C=m->C,V=m->V,F=m->F;
FBK(cublasCreate(&m->H));
/* opt #2: one resident device scratch slab, sized generously for this config's
* per-step forward+backward peak (B=1 in fm_forward/fm_backward). Transient
* buffers bump from here and are reclaimed en-masse each fm_forward/backward. */
{ size_t TC=(size_t)T*C; size_t scsz=(size_t)T*(size_t)m->nH*m->N*m->N;
size_t cap=(size_t)384*TC*sizeof(float)+(size_t)48*scsz*sizeof(float)+(32u<<20);
fmA_init(cap); }
/* rebuild the interleave (mirror rx_model_new_random) */
m->btype=(int*)calloc(m->n_blocks,sizeof(int)); m->bidx=(int*)calloc(m->n_blocks,sizeof(int));
if(m->n_attn==0){ for(int i=0;i<m->n_rwkv;i++){m->btype[i]=0;m->bidx[i]=i;} }
else{ int interval=m->n_rwkv/m->n_attn,pos=0;
for(int a=0;a<m->n_attn;a++){ int s=a*interval,e=(a+1)*interval;
for(int j=s;j<e;j++){m->btype[pos]=0;m->bidx[pos]=j;pos++;} m->btype[pos]=1;m->bidx[pos]=a;pos++; }
for(int j=m->n_attn*interval;j<m->n_rwkv;j++){m->btype[pos]=0;m->bidx[pos]=j;pos++;} }
/* globals */
{ int n; float*d,*g; const char*nm;
/* find emb/head/ln_out by name in the registry to get their host data */
m->emb=NULL; m->head=NULL; m->lnout_w=NULL; m->lnout_b=NULL;
for(int i=0;i<rx_ptmodel_num_params(cm);i++){ rx_ptmodel_param(cm,i,&nm,&d,&g,&n);
if(!strcmp(nm,"emb")) m->emb=fm_up(d,(size_t)V*C);
else if(!strcmp(nm,"head")) m->head=fm_up(d,(size_t)C*V);
else if(!strcmp(nm,"ln_out_w")) m->lnout_w=fm_up(d,C);
else if(!strcmp(nm,"ln_out_b")) m->lnout_b=fm_up(d,C); }
m->g_emb=fm_upg(m,"emb",V*C); m->g_head=fm_upg(m,"head",C*V);
m->g_lnout_w=fm_upg(m,"ln_out_w",C); m->g_lnout_b=fm_upg(m,"ln_out_b",C);
}
/* helper to fetch a host param pointer by name */
#define HOSTP(dst,nm_) do{ const char*_n; float*_d,*_g; int _k; dst=NULL; \
for(int _i=0;_i<rx_ptmodel_num_params(cm);_i++){ rx_ptmodel_param(cm,_i,&_n,&_d,&_g,&_k); \
if(!strcmp(_n,nm_)){ dst=_d; break; } } }while(0)
/* rwkv blocks */
m->rw=(fm_rwkv_blk*)calloc(m->n_rwkv>0?m->n_rwkv:1,sizeof(fm_rwkv_blk));
for(int i=0;i<m->n_rwkv;i++){
fm_rwkv_blk*B=&m->rw[i];
const rx_block_params*bp=rx_ptmodel_rwkv(cm,i);
B->tmP=bp->tm; B->cmP=bp->cm; /* host mirror (dims/layer_id + host ptrs) */
/* upload time-mix weights via the merged uploader */
if(tm_upload_weights(&bp->tm,&B->D)){ fprintf(stderr,"tm upload fail\n"); exit(1); }
/* time-mix grads: register each by the oracle name */
char nm[64];
#define TMG(field,cnt) do{ snprintf(nm,64,"rwkv%d." #field,i); B->Gd.field=fm_upg(m,nm,cnt); }while(0)
TMG(x_r,C);TMG(x_k,C);TMG(x_v,C);TMG(x_w,C);TMG(x_a,C);TMG(x_g,C);
TMG(receptance_w,C*C);TMG(key_w,C*C);TMG(value_w,C*C);TMG(output_w,C*C);
TMG(w0,C);TMG(w1,C*m->Dw);TMG(w2,m->Dw*C);
TMG(a0,C);TMG(a1,C*m->Da);TMG(a2,m->Da*C);
TMG(g1,C*m->Dg);TMG(g2,m->Dg*C);
if(i>0){ TMG(v0,C);TMG(v1,C*m->Dv);TMG(v2,m->Dv*C); } else { B->Gd.v0=B->Gd.v1=B->Gd.v2=NULL; }
TMG(k_k,C);TMG(k_a,C);TMG(r_k,C);TMG(ln_x_w,C);TMG(ln_x_b,C);
#undef TMG
/* channel-mix weights + grads */
float*hp;
HOSTP(hp,(snprintf(nm,64,"rwkv%d.ffn_x_k",i),nm)); B->cm.x_k=fm_up(hp,C);
HOSTP(hp,(snprintf(nm,64,"rwkv%d.ffn_key_w",i),nm)); B->cm.key_w=fm_up(hp,(size_t)C*F);
HOSTP(hp,(snprintf(nm,64,"rwkv%d.ffn_value_w",i),nm)); B->cm.value_w=fm_up(hp,(size_t)F*C);
snprintf(nm,64,"rwkv%d.ffn_x_k",i); B->cmG.x_k=fm_upg(m,nm,C);
snprintf(nm,64,"rwkv%d.ffn_key_w",i); B->cmG.key_w=fm_upg(m,nm,C*F);
snprintf(nm,64,"rwkv%d.ffn_value_w",i); B->cmG.value_w=fm_upg(m,nm,F*C);
/* ln1/ln2 */
HOSTP(hp,(snprintf(nm,64,"rwkv%d.ln1_w",i),nm)); B->ln1_w=fm_up(hp,C);
HOSTP(hp,(snprintf(nm,64,"rwkv%d.ln1_b",i),nm)); B->ln1_b=fm_up(hp,C);
HOSTP(hp,(snprintf(nm,64,"rwkv%d.ln2_w",i),nm)); B->ln2_w=fm_up(hp,C);
HOSTP(hp,(snprintf(nm,64,"rwkv%d.ln2_b",i),nm)); B->ln2_b=fm_up(hp,C);
snprintf(nm,64,"rwkv%d.ln1_w",i); B->g_ln1_w=fm_upg(m,nm,C);
snprintf(nm,64,"rwkv%d.ln1_b",i); B->g_ln1_b=fm_upg(m,nm,C);
snprintf(nm,64,"rwkv%d.ln2_w",i); B->g_ln2_w=fm_upg(m,nm,C);
snprintf(nm,64,"rwkv%d.ln2_b",i); B->g_ln2_b=fm_upg(m,nm,C);
/* point the host cm mirror at oracle host pointers (dims only used) */
B->cmP=bp->cm;
}
/* attn blocks */
m->at=(fm_attn_blk*)calloc(m->n_attn>0?m->n_attn:1,sizeof(fm_attn_blk));
for(int i=0;i<m->n_attn;i++){
fm_attn_blk*A=&m->at[i]; const rx_attn_params*ap=rx_ptmodel_attn(cm,i); A->P=*ap;
char nm[64]; float*hp;
HOSTP(hp,(snprintf(nm,64,"attn%d.ln1_w",i),nm)); A->ln1_w=fm_up(hp,C);
HOSTP(hp,(snprintf(nm,64,"attn%d.ln1_b",i),nm)); A->ln1_b=fm_up(hp,C);
HOSTP(hp,(snprintf(nm,64,"attn%d.receptance_w",i),nm)); A->Wq=fm_up(hp,(size_t)C*C);
HOSTP(hp,(snprintf(nm,64,"attn%d.key_w",i),nm)); A->Wk=fm_up(hp,(size_t)C*C);
HOSTP(hp,(snprintf(nm,64,"attn%d.value_w",i),nm)); A->Wv=fm_up(hp,(size_t)C*C);
HOSTP(hp,(snprintf(nm,64,"attn%d.output_w",i),nm)); A->Wo=fm_up(hp,(size_t)C*C);
HOSTP(hp,(snprintf(nm,64,"attn%d.ln2_w",i),nm)); A->ln2_w=fm_up(hp,C);
HOSTP(hp,(snprintf(nm,64,"attn%d.ln2_b",i),nm)); A->ln2_b=fm_up(hp,C);
HOSTP(hp,(snprintf(nm,64,"attn%d.ffn_x_k",i),nm)); A->cm.x_k=fm_up(hp,C);
HOSTP(hp,(snprintf(nm,64,"attn%d.ffn_key_w",i),nm)); A->cm.key_w=fm_up(hp,(size_t)C*F);
HOSTP(hp,(snprintf(nm,64,"attn%d.ffn_value_w",i),nm)); A->cm.value_w=fm_up(hp,(size_t)F*C);
/* grads (registered) */
snprintf(nm,64,"attn%d.ln1_w",i); fm_upg(m,nm,C); snprintf(nm,64,"attn%d.ln1_b",i); fm_upg(m,nm,C);
snprintf(nm,64,"attn%d.receptance_w",i); fm_upg(m,nm,C*C);
snprintf(nm,64,"attn%d.key_w",i); fm_upg(m,nm,C*C);
snprintf(nm,64,"attn%d.value_w",i); fm_upg(m,nm,C*C);
snprintf(nm,64,"attn%d.output_w",i); fm_upg(m,nm,C*C);
snprintf(nm,64,"attn%d.ln2_w",i); fm_upg(m,nm,C); snprintf(nm,64,"attn%d.ln2_b",i); fm_upg(m,nm,C);
snprintf(nm,64,"attn%d.ffn_x_k",i); fm_upg(m,nm,C);
snprintf(nm,64,"attn%d.ffn_key_w",i); fm_upg(m,nm,C*F);
snprintf(nm,64,"attn%d.ffn_value_w",i); fm_upg(m,nm,F*C);
}
#undef HOSTP
/* activations */
size_t TC=(size_t)T*C;
FCK(cudaMalloc(&m->tok,(size_t)T*sizeof(int)));
FCK(cudaMalloc(&m->dtg,(size_t)T*sizeof(int))); /* opt#3 resident targets */
FCK(cudaMalloc(&m->block_in,(size_t)(m->n_blocks+1)*TC*sizeof(float)));
FCK(cudaMalloc(&m->v_first,TC*sizeof(float)));
FCK(cudaMalloc(&m->lnout,TC*sizeof(float)));
FCK(cudaMalloc(&m->logits,(size_t)T*V*sizeof(float)));
FCK(cudaMalloc(&m->loss_t,(size_t)T*sizeof(double)));
return m;
}
static void fm_zero_grads(fm_model*m){
for(int i=0;i<m->ng;i++) FCK(cudaMemsetAsync(m->gptr[i],0,(size_t)m->gn[i]*sizeof(float),g_fm_stream));
}
/* ============================ channel-mix device (resident) ============================ */
/* forward: sh=shift(in); fk=sh@Wk; ka=relu(fk)^2; out=ka@Wv. (x_prev_init=NULL) */
static void fm_cm_fwd(fm_model*m,const fm_cmw*W,const rx_cm_params*P,const float*in,float*out,int T){
int C=P->c,F=P->ffn_dim; size_t BT=(size_t)T, tot=BT*C, totF=BT*F;
float*sh=(float*)fmA_alloc(tot*4),*fk=(float*)fmA_alloc(totF*4),*ka=(float*)fmA_alloc(totF*4);
k_cm_shift<<<fm_nblk(tot),256>>>(in,NULL,W->x_k,sh,1,T,C);
FBK(rm_gemm_cm(m->H,0,0,T,F,C,1.0f,sh,W->key_w,0.0f,fk));
k_sqrelu<<<fm_nblk(totF),256>>>(fk,ka,totF);
FBK(rm_gemm_cm(m->H,0,0,T,C,F,1.0f,ka,W->value_w,0.0f,out));
fmA_free(sh);fmA_free(fk);fmA_free(ka);
}
/* backward: dy -> dx + grads(G). recomputes sh,fk,ka. grads WRITTEN(=). */
static void fm_cm_bwd(fm_model*m,const fm_cmw*W,fm_cmw*G,const rx_cm_params*P,
const float*in,const float*dy,float*dx,int T){
int C=P->c,F=P->ffn_dim; size_t BT=(size_t)T,tot=BT*C,totF=BT*F;
float*sh=(float*)fmA_alloc(tot*4),*fk=(float*)fmA_alloc(totF*4),*ka=(float*)fmA_alloc(totF*4);
float*dka=(float*)fmA_alloc(totF*4),*dfk=(float*)fmA_alloc(totF*4),*dsh=(float*)fmA_alloc(tot*4);
k_cm_shift<<<fm_nblk(tot),256>>>(in,NULL,W->x_k,sh,1,T,C);
FBK(rm_gemm_cm(m->H,0,0,T,F,C,1.0f,sh,W->key_w,0.0f,fk));
k_sqrelu<<<fm_nblk(totF),256>>>(fk,ka,totF);
FBK(rm_gemm_cm(m->H,0,1,T,F,C,1.0f,dy,W->value_w,0.0f,dka)); /* dka=dy@Wv^T */
FBK(rm_gemm_cm(m->H,1,0,F,C,T,1.0f,ka,dy,0.0f,G->value_w)); /* G_Wv=ka^T@dy */
k_sqrelu_bwd<<<fm_nblk(totF),256>>>(fk,dka,dfk,totF);
FBK(rm_gemm_cm(m->H,0,1,T,C,F,1.0f,dfk,W->key_w,0.0f,dsh)); /* dsh=dfk@Wk^T */
FBK(rm_gemm_cm(m->H,1,0,C,F,T,1.0f,sh,dfk,0.0f,G->key_w)); /* G_Wk=sh^T@dfk */
FCK(cudaMemsetAsync(G->x_k,0,(size_t)C*4,g_fm_stream));
k_cm_shift_bwd<<<fm_nblk(tot),256>>>(in,NULL,dsh,W->x_k,G->x_k,dx,NULL,1,T,C);
fmA_free(sh);fmA_free(fk);fmA_free(ka);fmA_free(dka);fmA_free(dfk);fmA_free(dsh);
}
/* ============================ time-mix backward core (resident) ============================
* Device-resident port of timemix_backward_cuda's phase-2..4 + token-shift, on a
* PRE-FILLED forward cache c and resident weights D / grads Gd (zeroed by caller).
* x=h1, ddy=dx1, dvf=v_first values, ddx_out=dh1. xprev/d_xprev = NULL (B=1). */
static void fm_tm_bwd_core(fm_model*m,const rx_tm_params*P,const tm_dw*D,tm_dw*Gd,
tm_cache_d*c,const float*x,const float*ddy,const float*dvf,
float*ddx_out,const float*d_vf_in,float*d_vf_out,int T){
int C=P->c,nH=P->n_head,N=P->head_size,Dw=P->Dw,Da=P->Da,Dg=P->Dg,Dv=P->Dv,L=P->layer_id;
size_t BT=(size_t)T, tot=BT*C;
float *dtg,*dyb,*dg_post,*dy_wkv,*dr_post,*dkmod_post,*dvf_post;
float *gR,*gW,*gK,*gV,*gA,*gB,*dr_t,*dkmod_t,*dvf_t;
float *dkk,*dav,*dkvec,*dvvec,*dtk,*d_oa,*d_ow,*d_ov;
float *dsr,*dsk,*dsv,*dsw,*dsa,*dsg;
float *d_la,*d_tw,*d_lw,*d_sg,*d_lg,*d_lv=NULL,*d_vf_loc=NULL;
#define S(p,n) do{ p=(float*)fmA_alloc((size_t)(n)*sizeof(float)); }while(0)
S(dtg,tot);S(dyb,tot);S(dg_post,tot);S(dy_wkv,tot);S(dr_post,tot);S(dkmod_post,tot);S(dvf_post,tot);
S(gR,tot);S(gW,tot);S(gK,tot);S(gV,tot);S(gA,tot);S(gB,tot);
S(dr_t,tot);S(dkmod_t,tot);S(dvf_t,tot);
S(dkk,tot);S(dav,tot);S(dkvec,tot);S(dvvec,tot);S(dtk,tot);S(d_oa,tot);S(d_ow,tot);S(d_ov,tot);
S(dsr,tot);S(dsk,tot);S(dsv,tot);S(dsw,tot);S(dsa,tot);S(dsg,tot);
S(d_la,BT*Da);S(d_tw,BT*Dw);S(d_lw,BT*Dw);S(d_sg,BT*Dg);S(d_lg,BT*Dg);
if(Dv>0) S(d_lv,BT*Dv);
if(L>0){ S(d_vf_loc,tot); FCK(cudaMemsetAsync(d_vf_loc,0,tot*4,g_fm_stream)); }
#undef S
/* PHASE 2 */
FBK(rm_gemm_tm(m->H,0,1,T,C,C,1.0f,ddy,D->output_w,0.0f,dtg));
FBK(rm_gemm_tm(m->H,1,0,C,C,T,1.0f,c->tg,ddy,0.0f,Gd->output_w));
k_split<<<fm_nblk(tot),256>>>(dtg,c->gvec,c->yb,dyb,dg_post,tot);
k_bonus_bwd<<<fm_nblk(BT*nH),256>>>(dyb,c->V,c->R,c->K,D->r_k,c->bonus,dvf_post,dr_post,dkmod_post,Gd->r_k,nH,N,BT);
k_groupnorm_bwd<<<fm_nblk(BT*nH),256>>>(c->y,dyb,D->ln_x_w,c->gn_mean,c->gn_rstd,dy_wkv,Gd->ln_x_w,Gd->ln_x_b,nH,N,BT);
/* PHASE 3: WKV7 */
{ size_t stsz=(size_t)nH*N*N, sasz=(size_t)T*nH*N, scsz=(size_t)T*stsz;
float*SC=(float*)fmA_alloc(scsz*4); float*SA=(float*)fmA_alloc(sasz*4);
size_t shf=((size_t)N*N+7*N)*sizeof(float);
wkv7_fwd_kernel<<<nH,N,shf>>>(c->R,c->W,c->K,c->V,c->A,c->Bb,NULL,NULL,NULL,SC,SA,1,T,nH,N);
size_t shb=(2*(size_t)N*N+9*N)*sizeof(float);
wkv7_bwd_kernel<<<nH,N,shb>>>(c->R,c->W,c->K,c->V,c->A,c->Bb,dy_wkv,SC,SA,gR,gW,gK,gV,gA,gB,NULL,1,T,nH,N);
fmA_free(SC);fmA_free(SA); }
/* PHASE 4 */
k_p4_seed<<<fm_nblk(tot),256>>>(gR,dr_post,gK,dkmod_post,gV,dvf_post,dr_t,dkmod_t,dvf_t,tot);
k_comp_keymod<<<fm_nblk(tot),256>>>(gA,gB,c->avec,c->kk,c->kvec,D->k_a,dkmod_t,dkk,dav,dkvec,Gd->k_a,C,tot);
k_l2norm_bwd<<<fm_nblk(BT*nH),256>>>(c->tk,dkk,dtk,nH,N,BT);
k_tk_bwd<<<fm_nblk(tot),256>>>(dtk,D->k_k,c->kvec,dkvec,Gd->k_k,C,tot);
k_avec_bwd<<<fm_nblk(tot),256>>>(dav,c->avec,d_oa,Gd->a0,C,tot);
FBK(rm_gemm_tm(m->H,0,1,T,Da,C,1.0f,d_oa,D->a2,0.0f,d_la));
FBK(rm_gemm_tm(m->H,1,0,Da,C,T,1.0f,c->la,d_oa,0.0f,Gd->a2));
FBK(rm_gemm_tm(m->H,0,1,T,C,Da,1.0f,d_la,D->a1,0.0f,dsa));
FBK(rm_gemm_tm(m->H,1,0,C,Da,T,1.0f,c->sha,d_la,0.0f,Gd->a1));
k_wdecay_bwd<<<fm_nblk(tot),256>>>(gW,c->W,c->sigw,d_ow,Gd->w0,C,tot);
FBK(rm_gemm_tm(m->H,0,1,T,Dw,C,1.0f,d_ow,D->w2,0.0f,d_tw));
FBK(rm_gemm_tm(m->H,1,0,Dw,C,T,1.0f,c->tw,d_ow,0.0f,Gd->w2));
k_tanhderiv<<<fm_nblk(BT*Dw),256>>>(d_tw,c->tw,d_lw,BT*Dw);
FBK(rm_gemm_tm(m->H,0,1,T,C,Dw,1.0f,d_lw,D->w1,0.0f,dsw));
FBK(rm_gemm_tm(m->H,1,0,C,Dw,T,1.0f,c->shw,d_lw,0.0f,Gd->w1));
FBK(rm_gemm_tm(m->H,0,1,T,Dg,C,1.0f,dg_post,D->g2,0.0f,d_sg));
FBK(rm_gemm_tm(m->H,1,0,Dg,C,T,1.0f,c->sg,dg_post,0.0f,Gd->g2));
k_sigderiv<<<fm_nblk(BT*Dg),256>>>(d_sg,c->sg,d_lg,BT*Dg);
FBK(rm_gemm_tm(m->H,0,1,T,C,Dg,1.0f,d_lg,D->g1,0.0f,dsg));
FBK(rm_gemm_tm(m->H,1,0,C,Dg,T,1.0f,c->shg,d_lg,0.0f,Gd->g1));
if(L==0){ k_vvec0<<<fm_nblk(tot),256>>>(dvf_t,d_vf_in,dvvec,tot); }
else{
k_vvec1<<<fm_nblk(tot),256>>>(dvf_t,c->blend,dvf,c->vvec,D->v0,dvvec,d_vf_loc,d_ov,Gd->v0,C,tot);
FBK(rm_gemm_tm(m->H,0,1,T,Dv,C,1.0f,d_ov,D->v2,0.0f,d_lv));
FBK(rm_gemm_tm(m->H,1,0,Dv,C,T,1.0f,c->lv,d_ov,0.0f,Gd->v2));
FBK(rm_gemm_tm(m->H,0,1,T,C,Dv,1.0f,d_lv,D->v1,0.0f,dsv));
FBK(rm_gemm_tm(m->H,1,0,C,Dv,T,1.0f,c->shv,d_lv,0.0f,Gd->v1));
}
FBK(rm_gemm_tm(m->H,0,1,T,C,C,1.0f,dr_t,D->receptance_w,0.0f,dsr));
FBK(rm_gemm_tm(m->H,1,0,C,C,T,1.0f,c->shr,dr_t,0.0f,Gd->receptance_w));
FBK(rm_gemm_tm(m->H,0,1,T,C,C,1.0f,dkvec,D->key_w,0.0f,dsk));
FBK(rm_gemm_tm(m->H,1,0,C,C,T,1.0f,c->shk,dkvec,0.0f,Gd->key_w));
FBK(rm_gemm_tm(m->H,0,1,T,C,C,1.0f,dvvec,D->value_w,(L>0)?1.0f:0.0f,dsv));
FBK(rm_gemm_tm(m->H,1,0,C,C,T,1.0f,c->shv,dvvec,0.0f,Gd->value_w));
/* token-shift backward (dxprev=NULL) */
k_tokenshift_bwd<<<fm_nblk(tot),256>>>(x,NULL,dsr,dsk,dsv,dsw,dsa,dsg,
D->x_r,D->x_k,D->x_v,D->x_w,D->x_a,D->x_g,
Gd->x_r,Gd->x_k,Gd->x_v,Gd->x_w,Gd->x_a,Gd->x_g, ddx_out,NULL,1,T,C);
if(L>0 && d_vf_out) FCK(cudaMemcpyAsync(d_vf_out,d_vf_loc,tot*4,cudaMemcpyDeviceToDevice,g_fm_stream));
float*sc[]={dtg,dyb,dg_post,dy_wkv,dr_post,dkmod_post,dvf_post,gR,gW,gK,gV,gA,gB,
dr_t,dkmod_t,dvf_t,dkk,dav,dkvec,dvvec,dtk,d_oa,d_ow,d_ov,dsr,dsk,dsv,dsw,dsa,dsg,
d_la,d_tw,d_lw,d_sg,d_lg,d_lv,d_vf_loc};
for(unsigned i=0;i<sizeof(sc)/sizeof(sc[0]);i++) if(sc[i]) fmA_free(sc[i]);
}
/* ============================ RWKV block fwd/bwd ============================ */
static void fm_rwkv_fwd(fm_model*m,int idx,const float*xin,float*xout,int T){
fm_rwkv_blk*B=&m->rw[idx]; int C=m->C; size_t tot=(size_t)T*C;
float*h1=(float*)fmA_alloc(tot*4),*tm=(float*)fmA_alloc(tot*4),*h2=(float*)fmA_alloc(tot*4),*cm=(float*)fmA_alloc(tot*4);
k_ln_fwd<<<fm_nblk(T),256>>>(xin,B->ln1_w,B->ln1_b,h1,NULL,NULL,C,T);
tm_cache_d c; memset(&c,0,sizeof(c));
tm_fwd_device(&B->tmP,m->H,&B->D,h1,NULL,m->v_first,&c,tm,1,T); tmc_free(&c);
fm_k_add<<<fm_nblk(tot),256>>>(xin,tm,xout,tot); /* xout=x1 */
k_ln_fwd<<<fm_nblk(T),256>>>(xout,B->ln2_w,B->ln2_b,h2,NULL,NULL,C,T);
fm_cm_fwd(m,&B->cm,&B->cmP,h2,cm,T);
fm_k_addi<<<fm_nblk(tot),256>>>(xout,cm,tot); /* xout=x1+cm */
fmA_free(h1);fmA_free(tm);fmA_free(h2);fmA_free(cm);
}
static void fm_rwkv_bwd(fm_model*m,int idx,const float*xin,const float*dy,float*dx,
float*g_vf,int T){
fm_rwkv_blk*B=&m->rw[idx]; int C=m->C; size_t tot=(size_t)T*C;
float*h1=(float*)fmA_alloc(tot*4),*tm=(float*)fmA_alloc(tot*4),*x1=(float*)fmA_alloc(tot*4),*h2=(float*)fmA_alloc(tot*4);
float*dh2=(float*)fmA_alloc(tot*4),*dxa=(float*)fmA_alloc(tot*4),*dx1=(float*)fmA_alloc(tot*4),*dh1=(float*)fmA_alloc(tot*4);
/* recompute forward temporaries */
k_ln_fwd<<<fm_nblk(T),256>>>(xin,B->ln1_w,B->ln1_b,h1,NULL,NULL,C,T);
tm_cache_d c; memset(&c,0,sizeof(c));
tm_fwd_device(&B->tmP,m->H,&B->D,h1,NULL,m->v_first,&c,tm,1,T);
fm_k_add<<<fm_nblk(tot),256>>>(xin,tm,x1,tot);
k_ln_fwd<<<fm_nblk(T),256>>>(x1,B->ln2_w,B->ln2_b,h2,NULL,NULL,C,T);
/* backward */
fm_cm_bwd(m,&B->cm,&B->cmG,&B->cmP,h2,dy,dh2,T);
FCK(cudaMemsetAsync(B->g_ln2_w,0,(size_t)C*4,g_fm_stream)); FCK(cudaMemsetAsync(B->g_ln2_b,0,(size_t)C*4,g_fm_stream));
k_ln_bwd<<<fm_nblk(T),256>>>(x1,B->ln2_w,dh2,dxa,B->g_ln2_w,B->g_ln2_b,C,T);
fm_k_add<<<fm_nblk(tot),256>>>(dy,dxa,dx1,tot); /* dx1=dy+dxa */
/* time-mix bwd (uses filled cache c) */
if(B->tmP.layer_id==0)
fm_tm_bwd_core(m,&B->tmP,&B->D,&B->Gd,&c,h1,dx1,m->v_first,dh1,g_vf,NULL,T);
else{
float*d_vf_out=(float*)fmA_alloc(tot*4);
fm_tm_bwd_core(m,&B->tmP,&B->D,&B->Gd,&c,h1,dx1,m->v_first,dh1,NULL,d_vf_out,T);
fm_k_addi<<<fm_nblk(tot),256>>>(g_vf,d_vf_out,tot);
fmA_free(d_vf_out);
}
tmc_free(&c);
FCK(cudaMemsetAsync(B->g_ln1_w,0,(size_t)C*4,g_fm_stream)); FCK(cudaMemsetAsync(B->g_ln1_b,0,(size_t)C*4,g_fm_stream));
k_ln_bwd<<<fm_nblk(T),256>>>(xin,B->ln1_w,dh1,dxa,B->g_ln1_w,B->g_ln1_b,C,T);
fm_k_add<<<fm_nblk(tot),256>>>(dx1,dxa,dx,tot); /* dx=dx1+dxa */
fmA_free(h1);fmA_free(tm);fmA_free(x1);fmA_free(h2);fmA_free(dh2);fmA_free(dxa);fmA_free(dx1);fmA_free(dh1);
}
/* ============================ Attention block fwd/bwd ============================ */
/* moba core forward: packed q,k,v (device [T,C]) -> y (device [T,C]); leaves P. */
static void fm_attn_core_fwd(fm_model*m,fm_attn_blk*A,const float*q,const float*k,const float*v,
float*y,float*P,int T){
int C=m->C,nH=m->nH,N=m->N,cs=A->P.chunk_size<1?1:A->P.chunk_size,topk=A->P.moba_top_k;
float scale=1.0f/sqrtf((float)N); size_t tot=(size_t)T*C, nsc=(size_t)nH*T*T;
moba_dev d; moba_alloc(&d,tot,nsc);
float*Yh=(float*)fmA_alloc(tot*4);
int nb=moba_nblk(tot);
moba_k_pack<<<nb,256>>>(q,d.qh,1,T,nH,N);
moba_k_pack<<<nb,256>>>(k,d.kh,1,T,nH,N);
moba_k_pack<<<nb,256>>>(v,d.vh,1,T,nH,N);
moba_scores_softmax(m->H,&d,1,nH,T,N,cs,topk,scale);
for(int bh=0;bh<nH;bh++)
rm_gemm_moba(m->H,0,0,T,N,T,1.0f,d.P+(size_t)bh*T*T,d.vh+(size_t)bh*T*N,0.0f,Yh+(size_t)bh*T*N);
moba_k_unpack<<<nb,256>>>(Yh,y,1,T,nH,N);
if(P) FCK(cudaMemcpyAsync(P,d.P,nsc*4,cudaMemcpyDeviceToDevice,g_fm_stream));
moba_free(&d); fmA_free(Yh);
}
static void fm_attn_fwd(fm_model*m,int idx,const float*xin,float*xout,int T){
fm_attn_blk*A=&m->at[idx]; int C=m->C; size_t tot=(size_t)T*C;
float*h1=(float*)fmA_alloc(tot*4),*q=(float*)fmA_alloc(tot*4),*k=(float*)fmA_alloc(tot*4),*v=(float*)fmA_alloc(tot*4);
float*y=(float*)fmA_alloc(tot*4),*h2=(float*)fmA_alloc(tot*4),*cm=(float*)fmA_alloc(tot*4);
k_ln_fwd<<<fm_nblk(T),256>>>(xin,A->ln1_w,A->ln1_b,h1,NULL,NULL,C,T);
FBK(rm_gemm_tm(m->H,0,0,T,C,C,1.0f,h1,A->Wq,0.0f,q));
FBK(rm_gemm_tm(m->H,0,0,T,C,C,1.0f,h1,A->Wk,0.0f,k));
FBK(rm_gemm_tm(m->H,0,0,T,C,C,1.0f,h1,A->Wv,0.0f,v));
fm_attn_core_fwd(m,A,q,k,v,y,NULL,T);
FBK(rm_gemm_tm(m->H,0,0,T,C,C,1.0f,y,A->Wo,0.0f,xout)); /* att_out */
fm_k_addi<<<fm_nblk(tot),256>>>(xout,xin,tot); /* x1=x+att_out */
k_ln_fwd<<<fm_nblk(T),256>>>(xout,A->ln2_w,A->ln2_b,h2,NULL,NULL,C,T);
fm_cmw cmw={A->cm.x_k,A->cm.key_w,A->cm.value_w};
fm_cm_fwd(m,&cmw,&A->P.cm,h2,cm,T);
fm_k_addi<<<fm_nblk(tot),256>>>(xout,cm,tot);
fmA_free(h1);fmA_free(q);fmA_free(k);fmA_free(v);fmA_free(y);fmA_free(h2);fmA_free(cm);
}
static void fm_attn_bwd(fm_model*m,int idx,const float*xin,const float*dy,float*dx,int T){
fm_attn_blk*A=&m->at[idx]; int C=m->C,nH=m->nH,N=m->N;
int cs=A->P.chunk_size<1?1:A->P.chunk_size,topk=A->P.moba_top_k;
float scale=1.0f/sqrtf((float)N); size_t tot=(size_t)T*C, nsc=(size_t)nH*T*T;
char nm[64];
float*h1,*q,*k,*v,*y,*x1,*h2,*dh2,*dxa,*dx1,*d_attout,*dy_core,*dq,*dk,*dv,*dh1;
#define A_(p) do{ p=(float*)fmA_alloc(tot*4); }while(0)
A_(h1);A_(q);A_(k);A_(v);A_(y);A_(x1);A_(h2);A_(dh2);A_(dxa);A_(dx1);A_(d_attout);A_(dy_core);A_(dq);A_(dk);A_(dv);A_(dh1);
#undef A_
/* recompute forward */
k_ln_fwd<<<fm_nblk(T),256>>>(xin,A->ln1_w,A->ln1_b,h1,NULL,NULL,C,T);
FBK(rm_gemm_tm(m->H,0,0,T,C,C,1.0f,h1,A->Wq,0.0f,q));
FBK(rm_gemm_tm(m->H,0,0,T,C,C,1.0f,h1,A->Wk,0.0f,k));
FBK(rm_gemm_tm(m->H,0,0,T,C,C,1.0f,h1,A->Wv,0.0f,v));
fm_attn_core_fwd(m,A,q,k,v,y,NULL,T);
FBK(rm_gemm_tm(m->H,0,0,T,C,C,1.0f,y,A->Wo,0.0f,x1));
fm_k_addi<<<fm_nblk(tot),256>>>(x1,xin,tot); /* x1 */
k_ln_fwd<<<fm_nblk(T),256>>>(x1,A->ln2_w,A->ln2_b,h2,NULL,NULL,C,T);
/* --- backward --- */
fm_cmw cmw={A->cm.x_k,A->cm.key_w,A->cm.value_w};
int nF; float*gcx=fm_gfind(m,(snprintf(nm,64,"attn%d.ffn_x_k",idx),nm),&nF);
float*gck=fm_gfind(m,(snprintf(nm,64,"attn%d.ffn_key_w",idx),nm),&nF);
float*gcv=fm_gfind(m,(snprintf(nm,64,"attn%d.ffn_value_w",idx),nm),&nF);
fm_cmw cmg={gcx,gck,gcv};
fm_cm_bwd(m,&cmw,&cmg,&A->P.cm,h2,dy,dh2,T);
float*gl2w=fm_gfind(m,(snprintf(nm,64,"attn%d.ln2_w",idx),nm),&nF);
float*gl2b=fm_gfind(m,(snprintf(nm,64,"attn%d.ln2_b",idx),nm),&nF);
FCK(cudaMemsetAsync(gl2w,0,(size_t)C*4,g_fm_stream));FCK(cudaMemsetAsync(gl2b,0,(size_t)C*4,g_fm_stream));
k_ln_bwd<<<fm_nblk(T),256>>>(x1,A->ln2_w,dh2,dxa,gl2w,gl2b,C,T);
fm_k_add<<<fm_nblk(tot),256>>>(dy,dxa,dx1,tot); /* dx1=dy+dxa */
/* output proj: att_out=y@Wo -> dy_core=dx1@Wo^T ; G_Wo=y^T@dx1 */
float*gWo=fm_gfind(m,(snprintf(nm,64,"attn%d.output_w",idx),nm),&nF);
FBK(rm_gemm_tm(m->H,0,1,T,C,C,1.0f,dx1,A->Wo,0.0f,dy_core));
FBK(rm_gemm_tm(m->H,1,0,C,C,T,1.0f,y,dx1,0.0f,gWo));
/* attention core backward: dy_core -> dq,dk,dv (selection fixed) */
{
moba_dev d; moba_alloc(&d,tot,nsc);
float*dyh=(float*)fmA_alloc(tot*4),*dqh=(float*)fmA_alloc(tot*4),*dkh=(float*)fmA_alloc(tot*4),*dvh=(float*)fmA_alloc(tot*4);
float*dP=(float*)fmA_alloc(nsc*4),*dS=(float*)fmA_alloc(nsc*4);
int nb=moba_nblk(tot);
moba_k_pack<<<nb,256>>>(q,d.qh,1,T,nH,N);
moba_k_pack<<<nb,256>>>(k,d.kh,1,T,nH,N);
moba_k_pack<<<nb,256>>>(v,d.vh,1,T,nH,N);
moba_k_pack<<<nb,256>>>(dy_core,dyh,1,T,nH,N);
moba_scores_softmax(m->H,&d,1,nH,T,N,cs,topk,scale);
for(int bh=0;bh<nH;bh++){ size_t os=(size_t)bh*T*T,ov=(size_t)bh*T*N;
rm_gemm_moba(m->H,0,1,T,T,N,1.0f,dyh+ov,d.vh+ov,0.0f,dP+os);
rm_gemm_moba(m->H,1,0,T,N,T,1.0f,d.P+os,dyh+ov,0.0f,dvh+ov); }
moba_k_softmax_bwd<<<moba_nblk((size_t)nH*T),256>>>(d.P,dP,dS,(long)nH,T,scale);
for(int bh=0;bh<nH;bh++){ size_t os=(size_t)bh*T*T,ov=(size_t)bh*T*N;
rm_gemm_moba(m->H,0,0,T,N,T,1.0f,dS+os,d.kh+ov,0.0f,dqh+ov);
rm_gemm_moba(m->H,1,0,T,N,T,1.0f,dS+os,d.qh+ov,0.0f,dkh+ov); }
moba_k_unpack<<<nb,256>>>(dqh,dq,1,T,nH,N);
moba_k_unpack<<<nb,256>>>(dkh,dk,1,T,nH,N);
moba_k_unpack<<<nb,256>>>(dvh,dv,1,T,nH,N);
moba_free(&d);fmA_free(dyh);fmA_free(dqh);fmA_free(dkh);fmA_free(dvh);fmA_free(dP);fmA_free(dS);
}
/* q,k,v = h1@{Wq,Wk,Wv} -> dh1 (accumulate) + proj grads */
float*gWq=fm_gfind(m,(snprintf(nm,64,"attn%d.receptance_w",idx),nm),&nF);
float*gWk=fm_gfind(m,(snprintf(nm,64,"attn%d.key_w",idx),nm),&nF);
float*gWv=fm_gfind(m,(snprintf(nm,64,"attn%d.value_w",idx),nm),&nF);
FBK(rm_gemm_tm(m->H,0,1,T,C,C,1.0f,dq,A->Wq,0.0f,dh1));
FBK(rm_gemm_tm(m->H,1,0,C,C,T,1.0f,h1,dq,0.0f,gWq));
FBK(rm_gemm_tm(m->H,0,1,T,C,C,1.0f,dk,A->Wk,1.0f,dh1));
FBK(rm_gemm_tm(m->H,1,0,C,C,T,1.0f,h1,dk,0.0f,gWk));
FBK(rm_gemm_tm(m->H,0,1,T,C,C,1.0f,dv,A->Wv,1.0f,dh1));
FBK(rm_gemm_tm(m->H,1,0,C,C,T,1.0f,h1,dv,0.0f,gWv));
/* ln1: dh1 -> dxa + ln1 grads ; dx = dx1 + dxa */
float*gl1w=fm_gfind(m,(snprintf(nm,64,"attn%d.ln1_w",idx),nm),&nF);
float*gl1b=fm_gfind(m,(snprintf(nm,64,"attn%d.ln1_b",idx),nm),&nF);
FCK(cudaMemsetAsync(gl1w,0,(size_t)C*4,g_fm_stream));FCK(cudaMemsetAsync(gl1b,0,(size_t)C*4,g_fm_stream));
k_ln_bwd<<<fm_nblk(T),256>>>(xin,A->ln1_w,dh1,dxa,gl1w,gl1b,C,T);
fm_k_add<<<fm_nblk(tot),256>>>(dx1,dxa,dx,tot);
fmA_free(h1);fmA_free(q);fmA_free(k);fmA_free(v);fmA_free(y);fmA_free(x1);fmA_free(h2);
fmA_free(dh2);fmA_free(dxa);fmA_free(dx1);fmA_free(d_attout);fmA_free(dy_core);fmA_free(dq);fmA_free(dk);fmA_free(dv);fmA_free(dh1);
}
/* ============================ full fwd / bwd ============================ */
static void fm_forward(fm_model*m,const int*tokens,int T){
int C=m->C,V=m->V; size_t TC=(size_t)T*C;
fmA_reset(); /* opt #2: reclaim last pass's transient scratch */
/* opt#3: during capture the driver has already staged m->tok (resident);
* outside capture (eval / non-graph) copy the batch in as before. */
if(!g_fm_capturing) FCK(cudaMemcpyAsync(m->tok,tokens,(size_t)T*sizeof(int),cudaMemcpyHostToDevice,g_fm_stream));
float*x0=m->block_in;
fm_k_emb_fwd<<<fm_nblk(TC),256>>>(m->emb,m->tok,x0,T,C);
for(int pos=0;pos<m->n_blocks;pos++){
float*xin=m->block_in+(size_t)pos*TC, *xout=m->block_in+(size_t)(pos+1)*TC;
if(m->btype[pos]==0) fm_rwkv_fwd(m,m->bidx[pos],xin,xout,T);
else fm_attn_fwd(m,m->bidx[pos],xin,xout,T);
}
float*final=m->block_in+(size_t)m->n_blocks*TC;
k_ln_fwd<<<fm_nblk(T),256>>>(final,m->lnout_w,m->lnout_b,m->lnout,NULL,NULL,C,T);
FBK(rm_gemm_tm(m->H,0,0,T,V,C,1.0f,m->lnout,m->head,0.0f,m->logits));
fm_sync();
}
static double fm_backward(fm_model*m,const int*targets,int T,int ignore_index){
int C=m->C,V=m->V; size_t TC=(size_t)T*C;
fmA_reset(); /* opt #2: reclaim the forward pass's transient scratch */
fm_zero_grads(m);
int nv=0; for(int t=0;t<T;t++) if(!(ignore_index>=0 && targets[t]==ignore_index)) nv++;
/* opt#3: targets live in resident m->dtg (stable ptr); driver stages them
* before graph replay, so skip the copy during capture. */
if(!g_fm_capturing) FCK(cudaMemcpyAsync(m->dtg,targets,(size_t)T*sizeof(int),cudaMemcpyHostToDevice,g_fm_stream));
float*dlogits=(float*)fmA_alloc((size_t)T*V*sizeof(float));
fm_k_ce_bwd<<<fm_nblk(T),256>>>(m->logits,m->dtg,dlogits,m->loss_t,T,V,ignore_index,nv);
fm_sync();
/* loss = mean over valid. Under capture the per-token loss stays in the
* resident m->loss_t buffer; the driver reduces it after replay (a D2H copy
* here would abort the capture), so return 0 as a placeholder. */
double loss=0.0;
if(!g_fm_capturing){
double*lh=(double*)malloc((size_t)T*sizeof(double));
FCK(cudaMemcpy(lh,m->loss_t,(size_t)T*sizeof(double),cudaMemcpyDeviceToHost));
for(int t=0;t<T;t++) loss+=lh[t]; loss = nv>0? loss/(double)nv : 0.0; free(lh);
}
/* head: dlogits -> dlnout + g_head */
float*dlnout=(float*)fmA_alloc(TC*4),*dfinal=(float*)fmA_alloc(TC*4);
float*dcur=(float*)fmA_alloc(TC*4),*dprev=(float*)fmA_alloc(TC*4),*g_vf=(float*)fmA_alloc(TC*4);
FCK(cudaMemsetAsync(g_vf,0,TC*4,g_fm_stream));
FBK(rm_gemm_tm(m->H,1,0,C,V,T,1.0f,m->lnout,dlogits,0.0f,m->g_head)); /* g_head=lnout^T@dlogits */
FBK(rm_gemm_tm(m->H,0,1,T,C,V,1.0f,dlogits,m->head,0.0f,dlnout)); /* dlnout=dlogits@head^T */
/* ln_out bwd */
float*final=m->block_in+(size_t)m->n_blocks*TC;
k_ln_bwd<<<fm_nblk(T),256>>>(final,m->lnout_w,dlnout,dfinal,m->g_lnout_w,m->g_lnout_b,C,T);
/* reverse blocks */
FCK(cudaMemcpyAsync(dcur,dfinal,TC*4,cudaMemcpyDeviceToDevice,g_fm_stream));
for(int pos=m->n_blocks-1;pos>=0;pos--){
float*xin=m->block_in+(size_t)pos*TC;
if(m->btype[pos]==0) fm_rwkv_bwd(m,m->bidx[pos],xin,dcur,dprev,g_vf,T);
else fm_attn_bwd(m,m->bidx[pos],xin,dcur,dprev,T);
float*tmp=dcur;dcur=dprev;dprev=tmp;
}
/* embedding scatter-add: dcur holds grad on emb output */
FCK(cudaMemsetAsync(m->g_emb,0,(size_t)V*C*sizeof(float),g_fm_stream));
fm_k_emb_bwd<<<fm_nblk(TC),256>>>(m->tok,dcur,m->g_emb,T,C);
fm_sync();
fmA_free(dlogits);fmA_free(dlnout);fmA_free(dfinal);fmA_free(dcur);fmA_free(dprev);fmA_free(g_vf);
return loss;
}
Reference in New Issue View Git Blame Copy Permalink