/* 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;
}