/* SPDX-License-Identifier: AGPL-3.0-or-later */
/* Copyright (C) 2026 David (Corey) Herrera */
/*
 * train_gpu.cu — On-device GPU training driver for the RWKV-x SR4-vs-AdamW scale
 * test (Phase 3, Chunk B).  The GPU counterpart of the scale test's CPU oracle
 * training driver: same enwik8 byte data pipeline, same batch
 * sampling, same next-byte train loop, same eval + logging (identical CSV), but
 * the model forward/backward and the optimizer step run ON THE K80, with weights
 * + optimizer state RESIDENT on-device (uploaded once).
 *
 * REUSE (do not rebuild): cuda_kernels/fullmodel.cu gives the device-resident
 * full RWKV-x fwd+bwd (fm_build / fm_forward / fm_backward, parity-green vs the
 * CPU oracle) and its grad registry (fm_gfind).  cuda_kernels/optim.cu gives the
 * GPU SR4 step (fc_sr4_scale + fc_sr4_update, bit-exact vs CPU) and the AdamW
 * step (fc_adamw_update).  This file adds ONLY: the data/loop/eval/logging glue,
 * the SR4 <-> continuous split for the ternary projections done on-device, the
 * per-column scale AdamW, and B>1 gradient accumulation.  Nothing in the merged
 * component files is edited.
 *
 * Optimizer split (identical policy to src/train.c / test_overfit.c):
 *   SR4 arm   : the big linear projections (attn + time-mix receptance/key/value/
 *               output, channel-mix ffn_key/ffn_value) are ternary {-1,0,+1} from
 *               an int4 accumulator q with a per-OUTPUT-COLUMN fp32 scale;
 *               W_eff[r,c] = scale[c]*ternary(q)[r,c].  Everything else (emb, head,
 *               LN, biases, lerps, LoRA gates, AND the scales) is continuous AdamW.
 *   AdamW arm : every parameter continuous, AdamW.
 *
 * Acceptance: run this and src/train.c (CPU) from the IDENTICAL seed/init/data/
 * batch order and compare per-step training loss.  Because the SR4 step is
 * bit-exact and fwd+bwd is parity-green, the SR4 arm tracks the CPU loss curve
 * very tightly (small fp drift only).
 *
 * Build (K80 host, GPU 0, from foundry root):
 *   gcc-11 -O2 -I. -c rwkvx/rx_pretrain.c        -o rx_pretrain_tg.o
 *   gcc-11 -O2 -I. -c data/fd_byte_dataset.c     -o fd_byte_dataset_tg.o
 *   gcc-11 -O2 -I. -c ternary_optim/ternary_optim.c -o ternary_optim_tg.o
 *   nvcc -O2 -arch=sm_37 -ccbin gcc-11 -isystem cuda_kernels/nvcc_shim -I. \
 *        cuda_kernels/train_gpu.cu rx_pretrain_tg.o fd_byte_dataset_tg.o \
 *        ternary_optim_tg.o -lcublas -lm -o build/train_gpu
 */
#include <cuda_runtime.h>
#include <cublas_v2.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include <stdint.h>
#include <time.h>

#include "cuda_kernels/fullmodel.cu"   /* fm_* device model + rx_pretrain.h + block kernels */
#include "cuda_kernels/optim.cu"       /* fc_sr4_scale / fc_sr4_update / fc_adamw_update */

#include "ternary_optim/ternary_optim.h"   /* host q-init (identical to CPU) */
#include "data/fd_byte_dataset.h"

#define LN2 0.69314718055994530942
#define SR4_Q_CLAMP 7
#define SR4_THETA   2

#define TGK(call) do{ cudaError_t e_=(call); if(e_!=cudaSuccess){ \
  fprintf(stderr,"CUDA %s:%d %s\n",__FILE__,__LINE__,cudaGetErrorString(e_)); exit(1);} }while(0)

static inline unsigned tg_nb(size_t n){ return (unsigned)((n+255)/256); }

/* ================================ device glue kernels ================================ */
/* ternarize the int4 accumulator: w = 0 if |q|<theta else sign(q).  (== ternary_optim_materialize) */
__global__ void tg_tern_from_q(const int8_t*q,float*tern,int theta,int numel){
    int i=blockIdx.x*blockDim.x+threadIdx.x; if(i>=numel)return;
    int v=q[i]; tern[i]=(v>=theta)?1.0f:(v<=-theta)?-1.0f:0.0f;
}
/* W_eff[r,c] = scale[c]*tern[r,c]  (row-major [rows,cols]) */
__global__ void tg_apply_scale(float*Weff,const float*tern,const float*scale,int rows,int cols){
    size_t g=(size_t)blockIdx.x*blockDim.x+threadIdx.x; size_t n=(size_t)rows*cols; if(g>=n)return;
    int c=(int)(g%(size_t)cols); Weff[g]=scale[c]*tern[g];
}
/* gq[r,c] = grad[r,c]*scale[c]  (dL/d ternaryW) */
__global__ void tg_gq(const float*grad,const float*scale,float*gq,int rows,int cols){
    size_t g=(size_t)blockIdx.x*blockDim.x+threadIdx.x; size_t n=(size_t)rows*cols; if(g>=n)return;
    int c=(int)(g%(size_t)cols); gq[g]=grad[g]*scale[c];
}
/* per-column scale grad: sgrad[c] = sum_r grad[r,c]*tern[r,c], double-accumulated
 * in the SAME index order as the CPU (r ascending) to minimize fp drift. */
__global__ void tg_scale_grad(const float*grad,const float*tern,float*sgrad,int rows,int cols){
    int c=blockIdx.x*blockDim.x+threadIdx.x; if(c>=cols)return;
    double s=0.0;
    for(int r=0;r<rows;r++){ size_t k=(size_t)r*cols+c; s+=(double)grad[k]*(double)tern[k]; }
    sgrad[c]=(float)s;
}
__global__ void tg_scale_vec(float*v,float k,int n){ int i=blockIdx.x*blockDim.x+threadIdx.x; if(i<n)v[i]*=k; }
/* opt#3 graph variant: the scale factor (scale_lr/lr, varies with the cosine LR)
 * is read from the device optimizer-scalar buffer opt[3]. */
__global__ void tg_scale_vec_g(float*v,const float*opt,int n){ int i=blockIdx.x*blockDim.x+threadIdx.x; if(i<n)v[i]*=opt[3]; }
__global__ void tg_acc_add(float*acc,const float*g,size_t n){ size_t i=(size_t)blockIdx.x*blockDim.x+threadIdx.x; if(i<n)acc[i]+=g[i]; }
__global__ void tg_scale_write(float*dst,const float*acc,float inv,size_t n){ size_t i=(size_t)blockIdx.x*blockDim.x+threadIdx.x; if(i<n)dst[i]=acc[i]*inv; }
__global__ void tg_fill(float*v,float val,int n){ int i=blockIdx.x*blockDim.x+threadIdx.x; if(i<n)v[i]=val; }

/* ================================ resident optimizer launchers ================================ */
static void adamw_dev(float*data,const float*grad,float*mm,float*vv,int64_t n,
                      int t,float lr,float b1,float b2,float eps,float wd,int is_adamw){
    if(n<=0) return;
    float bc1=1.0f-powf(b1,(float)t), bc2=1.0f-powf(b2,(float)t);
    int th=256; int64_t bl=(n+th-1)/th;
    fc_adamw_update<<<(unsigned)bl,th>>>(data,grad,mm,vv,n,lr,b1,b2,eps,wd,bc1,bc2,is_adamw);
}
/* SR4 q-step on device-resident pointers (batch_size==1: grad already averaged). */
static void sr4_dev(int8_t*q,const float*gq,int numel,uint64_t counter,float target_step,
                    uint64_t seed,int64_t idx_base,int q_clamp,int use_dither,float*d_s){
    float inv_bs=1.0f;
    fc_sr4_scale<<<1,1>>>(gq,numel,inv_bs,target_step,d_s);
    int th=256,bl=(numel+th-1)/th;
    fc_sr4_update<<<bl,th>>>(q,gq,numel,inv_bs,d_s,seed,idx_base,counter,q_clamp,use_dither);
}

/* opt#3 graph variants: per-step scalars come from device buffers (d_opt / d_ctr)
 * so the captured kernel nodes replay with fresh values while pointers stay fixed. */
static void adamw_dev_g(float*data,const float*grad,float*mm,float*vv,int64_t n,
                        const float*d_opt,float b1,float b2,float eps,float wd,int is_adamw){
    if(n<=0) return;
    int th=256; int64_t bl=(n+th-1)/th;
    fc_adamw_update_g<<<(unsigned)bl,th>>>(data,grad,mm,vv,n,d_opt,b1,b2,eps,wd,is_adamw);
}
static void sr4_dev_g(int8_t*q,const float*gq,int numel,const unsigned long long*d_ctr,
                      float target_step,uint64_t seed,int64_t idx_base,int q_clamp,int use_dither,float*d_s){
    float inv_bs=1.0f;
    fc_sr4_scale<<<1,1>>>(gq,numel,inv_bs,target_step,d_s);
    int th=256,bl=(numel+th-1)/th;
    fc_sr4_update_g<<<bl,th>>>(q,gq,numel,inv_bs,d_s,seed,idx_base,d_ctr,q_clamp,use_dither);
}

/* ================================ param classification (== train.c) ================================ */
static const char *suffix_after_dot(const char *name){ const char*d=strrchr(name,'.'); return d?d+1:name; }
static int is_ternary_name(const char *name){
    const char *s=suffix_after_dot(name);
    return !strcmp(s,"receptance_w")||!strcmp(s,"key_w")||!strcmp(s,"value_w")||
           !strcmp(s,"output_w")||!strcmp(s,"ffn_key_w")||!strcmp(s,"ffn_value_w");
}
static int rows_of(const char *name,int numel,int C,int ffn){
    const char*s=suffix_after_dot(name); (void)numel;
    if(!strcmp(s,"ffn_value_w")) return ffn; return C;
}

/* device weight pointer for an oracle param name (mirror of fm_build registration). */
static float* fm_wfind(fm_model*m,const char*name){
    if(!strcmp(name,"emb")) return m->emb;
    if(!strcmp(name,"head")) return m->head;
    if(!strcmp(name,"ln_out_w")) return m->lnout_w;
    if(!strcmp(name,"ln_out_b")) return m->lnout_b;
    int idx=-1; char suf[48]={0};
    if(sscanf(name,"rwkv%d.%47s",&idx,suf)==2 && idx>=0 && idx<m->n_rwkv){
        fm_rwkv_blk*B=&m->rw[idx]; tm_dw*D=&B->D;
        #define M(s,p) if(!strcmp(suf,s)) return p
        M("x_r",D->x_r);M("x_k",D->x_k);M("x_v",D->x_v);M("x_w",D->x_w);M("x_a",D->x_a);M("x_g",D->x_g);
        M("receptance_w",D->receptance_w);M("key_w",D->key_w);M("value_w",D->value_w);M("output_w",D->output_w);
        M("w0",D->w0);M("w1",D->w1);M("w2",D->w2);
        M("a0",D->a0);M("a1",D->a1);M("a2",D->a2);
        M("g1",D->g1);M("g2",D->g2);
        M("v0",D->v0);M("v1",D->v1);M("v2",D->v2);
        M("k_k",D->k_k);M("k_a",D->k_a);M("r_k",D->r_k);M("ln_x_w",D->ln_x_w);M("ln_x_b",D->ln_x_b);
        M("ffn_x_k",B->cm.x_k);M("ffn_key_w",B->cm.key_w);M("ffn_value_w",B->cm.value_w);
        M("ln1_w",B->ln1_w);M("ln1_b",B->ln1_b);M("ln2_w",B->ln2_w);M("ln2_b",B->ln2_b);
        #undef M
    }
    idx=-1; memset(suf,0,sizeof suf);
    if(sscanf(name,"attn%d.%47s",&idx,suf)==2 && idx>=0 && idx<m->n_attn){
        fm_attn_blk*A=&m->at[idx];
        #define M(s,p) if(!strcmp(suf,s)) return p
        M("ln1_w",A->ln1_w);M("ln1_b",A->ln1_b);M("ln2_w",A->ln2_w);M("ln2_b",A->ln2_b);
        M("receptance_w",A->Wq);M("key_w",A->Wk);M("value_w",A->Wv);M("output_w",A->Wo);
        M("ffn_x_k",A->cm.x_k);M("ffn_key_w",A->cm.key_w);M("ffn_value_w",A->cm.value_w);
        #undef M
    }
    return NULL;
}

/* ================================ config (== train.c) ================================ */
typedef struct {
    int C,n_head,head_size,n_rwkv,n_attn,ffn,lora;
    int seq_len,batch,steps,eval_every,eval_batches;
    uint64_t seed;
    int use_sr4;
    float target_step;
    float lr0,lr_min;
    float scale_lr;
    const char *logpath,*datapath;
} cfg_t;

static cfg_t default_cfg(void){
    cfg_t c; memset(&c,0,sizeof c);
    c.C=128; c.n_head=4; c.head_size=32; c.n_rwkv=3; c.n_attn=1;
    c.ffn=512; c.lora=32;
    c.seq_len=512; c.batch=1; c.steps=1000; c.eval_every=5; c.eval_batches=8;
    c.seed=1234ULL;
    c.use_sr4=0; c.target_step=0.003f;
    c.lr0=1e-3f; c.lr_min=1e-4f; c.scale_lr=1e-2f;
    c.logpath=NULL; c.datapath="data/enwik8";
    return c;
}
static rx_pt_config to_pt_config(const cfg_t *c){
    rx_pt_config pc; memset(&pc,0,sizeof pc);
    pc.n_embd=c->C; pc.n_head=c->n_head; pc.head_size=c->head_size;
    pc.n_rwkv_layer=c->n_rwkv; pc.n_attn_layer=c->n_attn;
    pc.vocab_size=FD_BYTE_VOCAB;
    pc.Dw=c->lora; pc.Da=c->lora; pc.Dg=c->lora; pc.Dv=c->lora;
    pc.ffn_dim=c->ffn;
    pc.chunk_size=c->seq_len; pc.moba_top_k=0;
    return pc;
}
static float cosine_lr(float lr0,float lr_min,int step,int total){
    if(total<=0||step>=total) return lr_min;
    float cs=0.5f*(1.0f+cosf((float)M_PI*(float)step/(float)total));
    return lr_min+(lr0-lr_min)*cs;
}
static double now_sec(void){ struct timespec ts; clock_gettime(CLOCK_MONOTONIC,&ts);
    return (double)ts.tv_sec+1e-9*(double)ts.tv_nsec; }
static void iso_now(char*buf,size_t n){ time_t tt=time(NULL); struct tm tmv; localtime_r(&tt,&tmv);
    strftime(buf,n,"%Y-%m-%dT%H:%M:%S",&tmv); }

/* L2 over ALL model param grads (device registry) — matches train.c grad_l2. */
static double grad_l2_gpu(fm_model*m,float*hbuf){
    double s=0.0;
    for(int i=0;i<m->ng;i++){ int n=m->gn[i];
        TGK(cudaMemcpy(hbuf,m->gptr[i],(size_t)n*sizeof(float),cudaMemcpyDeviceToHost));
        for(int k=0;k<n;k++) s+=(double)hbuf[k]*(double)hbuf[k]; }
    return sqrt(s);
}

/* held-out val: mean bpb + top-1 acc (identical protocol to train.c eval_metrics). */
static double eval_metrics_gpu(fm_model*m,const fd_byte_split*val,uint64_t seed,int seq_len,
                               int n_batches,int*in,int*tg,float*hlog,double*out_acc){
    double sum_ce=0.0; int n=0; long correct=0,total=0; int V=m->V;
    for(int b=0;b<n_batches;b++){
        fd_byte_get_batch(val,seed,0,(uint64_t)b,0,1,1,in,tg);
        fm_forward(m,in,seq_len);
        TGK(cudaMemcpy(hlog,m->logits,(size_t)seq_len*V*sizeof(float),cudaMemcpyDeviceToHost));
        double ce=rx_cross_entropy_forward(hlog,tg,seq_len,V,-1);
        sum_ce+=ce; n++;
        for(int t=0;t<seq_len;t++){ int tgt=tg[t]; if(tgt<0)continue;
            const float*row=&hlog[(size_t)t*V]; int am=0; float best=row[0];
            for(int v=1;v<V;v++) if(row[v]>best){best=row[v];am=v;}
            if(am==tgt)correct++; total++; }
    }
    if(out_acc)*out_acc= total?(double)correct/(double)total:0.0;
    return (n?sum_ce/n:0.0)/LN2;
}

/* ================================ ternary-matrix device bookkeeping ================================ */
typedef struct {
    char name[48];
    float *Weff,*grad;              /* device W_eff buffer + dL/dW_eff (fm registry) */
    int rows,cols,numel;
    int64_t idx_base;
    int8_t *d_q;                    /* int4 accumulator (device) */
    float *d_scale,*d_ms,*d_vs;     /* per-column scale + AdamW state [cols] */
    float *d_tern,*d_gq;            /* scratch [numel] */
    float *d_sgrad,*d_s;            /* scratch: scale grad [cols], step-scale scalar */
    float *snap;                    /* host prev ternary classification (churn) */
    int8_t *hq;                     /* host q scratch (panel) */
} tern_mat;

typedef struct { float*data,*grad,*mm,*vv; int64_t n; } cont_p;

typedef struct { double churn,pct_neg,pct_zero,pct_pos,rail_pct,dead_pct,scale_mean,scale_min,scale_max; } sr4_panel;

/* compute the SR4 health panel from the POST-step device q + scale (host reduction). */
static void sr4_panel_compute(tern_mat*tm,int n_tern,int snap_valid,sr4_panel*p,float*hscale_max_buf){
    (void)hscale_max_buf;
    long tot=0,cneg=0,czero=0,cpos=0,crail=0,cdead=0,cchurn=0;
    double ssum=0.0,smin=1e30,smax=-1e30; long scols=0;
    for(int t=0;t<n_tern;t++){
        tern_mat*tt=&tm[t]; int n=tt->numel;
        TGK(cudaMemcpy(tt->hq,tt->d_q,(size_t)n,cudaMemcpyDeviceToHost));
        float*hs=(float*)malloc((size_t)tt->cols*sizeof(float));
        TGK(cudaMemcpy(hs,tt->d_scale,(size_t)tt->cols*sizeof(float),cudaMemcpyDeviceToHost));
        for(int k=0;k<n;k++){
            int q=tt->hq[k];
            float v=(q>=SR4_THETA)?1.0f:(q<=-SR4_THETA)?-1.0f:0.0f;
            if(v<0.0f)cneg++; else if(v>0.0f)cpos++; else czero++;
            int aq=q<0?-q:q;
            if(aq==SR4_Q_CLAMP)crail++;
            if(aq<SR4_THETA)cdead++;
            if(snap_valid && tt->snap[k]!=v)cchurn++;
            tt->snap[k]=v; tot++;
        }
        for(int cc=0;cc<tt->cols;cc++){ float s=hs[cc]; ssum+=s; if(s<smin)smin=s; if(s>smax)smax=s; scols++; }
        free(hs);
    }
    double dt=tot?(double)tot:1.0;
    p->churn=snap_valid?(double)cchurn/dt:0.0;
    p->pct_neg=(double)cneg/dt; p->pct_zero=(double)czero/dt; p->pct_pos=(double)cpos/dt;
    p->rail_pct=(double)crail/dt; p->dead_pct=(double)cdead/dt;
    p->scale_mean=scols?ssum/(double)scols:0.0; p->scale_min=scols?smin:0.0; p->scale_max=scols?smax:0.0;
}

/* ================================================================== */
int main(int argc,char**argv){
    setvbuf(stdout,NULL,_IOLBF,0);   /* line-buffered for live tails */
    cfg_t c=default_cfg();
    for(int i=1;i<argc;i++){
        const char*a=argv[i];
        #define NEXT (++i<argc?argv[i]:"")
        if      (!strcmp(a,"--opt")){ const char*v=NEXT; c.use_sr4=!strcmp(v,"sr4"); }
        else if (!strcmp(a,"--target-step")) c.target_step=(float)atof(NEXT);
        else if (!strcmp(a,"--lr"))          c.lr0=(float)atof(NEXT);
        else if (!strcmp(a,"--lr-min"))      c.lr_min=(float)atof(NEXT);
        else if (!strcmp(a,"--scale-lr"))    c.scale_lr=(float)atof(NEXT);
        else if (!strcmp(a,"--seq-len"))     c.seq_len=atoi(NEXT);
        else if (!strcmp(a,"--batch"))       c.batch=atoi(NEXT);
        else if (!strcmp(a,"--steps"))       c.steps=atoi(NEXT);
        else if (!strcmp(a,"--eval-every"))  c.eval_every=atoi(NEXT);
        else if (!strcmp(a,"--eval-batches"))c.eval_batches=atoi(NEXT);
        else if (!strcmp(a,"--seed"))        c.seed=(uint64_t)strtoull(NEXT,NULL,10);
        else if (!strcmp(a,"--C"))           c.C=atoi(NEXT);
        else if (!strcmp(a,"--n-head"))      c.n_head=atoi(NEXT);
        else if (!strcmp(a,"--head-size"))   c.head_size=atoi(NEXT);
        else if (!strcmp(a,"--n-rwkv"))      c.n_rwkv=atoi(NEXT);
        else if (!strcmp(a,"--n-attn"))      c.n_attn=atoi(NEXT);
        else if (!strcmp(a,"--ffn"))         c.ffn=atoi(NEXT);
        else if (!strcmp(a,"--lora"))        c.lora=atoi(NEXT);
        else if (!strcmp(a,"--data"))        c.datapath=NEXT;
        else if (!strcmp(a,"--log"))         c.logpath=NEXT;
        else if (!strcmp(a,"--device"))      { (void)NEXT; }   /* accepted; this binary is gpu */
        else { fprintf(stderr,"unknown arg: %s\n",a); return 2; }
        #undef NEXT
    }
    if(c.seed==FD_BYTE_SEED_UNSET){ fprintf(stderr,"seed must be nonzero\n"); return 2; }
    if(c.head_size*c.n_head!=c.C){ fprintf(stderr,"head_size*n_head (%d) must == C (%d)\n",c.head_size*c.n_head,c.C); return 2; }
    if(c.eval_every<1) c.eval_every=1;

    int dev=0; cudaSetDevice(dev); cudaDeviceProp pr; cudaGetDeviceProperties(&pr,dev);

    FILE*logfp=c.logpath?fopen(c.logpath,"w"):NULL;
    #define LOGF(...) do{ printf(__VA_ARGS__); if(logfp){fprintf(logfp,__VA_ARGS__); fflush(logfp);} }while(0)

    /* ---- data ---- */
    fd_byte_corpus corpus;
    if(fd_byte_corpus_open(&corpus,c.datapath)!=0){ fprintf(stderr,"failed to open corpus '%s'\n",c.datapath); return 1; }
    fd_byte_split train_s,val_s;
    if(fd_byte_split_enwik8_train(&train_s,&corpus,(size_t)c.seq_len)!=0 ||
       fd_byte_split_enwik8_val  (&val_s,  &corpus,(size_t)c.seq_len)!=0){
        fprintf(stderr,"failed to build enwik8 splits\n"); return 1; }

    /* ---- model (CPU oracle struct, uploaded to device once) ---- */
    rx_pt_config pc=to_pt_config(&c);
    rx_ptmodel*cm=rx_model_new_random(&pc,c.seed);
    if(!cm){ fprintf(stderr,"model init failed\n"); return 1; }
    fm_model*m=fm_build(cm,c.seq_len);       /* upload ALL weights ONCE (resident) */
    int nP=rx_ptmodel_num_params(cm);

    long total_params=0,ternary_params=0;
    for(int i=0;i<nP;i++){ const char*nm; float*d,*g; int n; rx_ptmodel_param(cm,i,&nm,&d,&g,&n);
        total_params+=n; if(is_ternary_name(nm)) ternary_params+=n; }

    /* ---- optimizer state (device-resident) ---- */
    tern_mat*tmat=(tern_mat*)calloc(nP,sizeof(tern_mat)); int n_tern=0;
    cont_p *cont=(cont_p*)calloc(nP,sizeof(cont_p)); int n_cont=0;
    int64_t idx_base=0;
    for(int i=0;i<nP;i++){
        const char*nm; float*d,*g; int n; rx_ptmodel_param(cm,i,&nm,&d,&g,&n);
        float*wdev=fm_wfind(m,nm);
        int gn; float*gdev=fm_gfind(m,nm,&gn);
        if(!wdev||!gdev||gn!=n){ fprintf(stderr,"param map fail: %s (w=%p g=%p gn=%d n=%d)\n",nm,(void*)wdev,(void*)gdev,gn,n); return 1; }
        if(c.use_sr4 && is_ternary_name(nm)){
            tern_mat*tt=&tmat[n_tern++];
            snprintf(tt->name,48,"%s",nm);
            tt->Weff=wdev; tt->grad=gdev; tt->numel=n;
            tt->rows=rows_of(nm,n,c.C,c.ffn); tt->cols=n/tt->rows;
            tt->idx_base=idx_base; idx_base+=n;
            /* host ternary_optim purely to reproduce the CPU q-init bit-for-bit */
            ternary_optim_config oc; ternary_optim_config_default(&oc);
            oc.seed=42; oc.idx_base=tt->idx_base;
            ternary_optim*ho=ternary_optim_create(tt->rows,tt->cols,&oc);
            const int8_t*qh=ternary_optim_q(ho);
            TGK(cudaMalloc(&tt->d_q,(size_t)n)); TGK(cudaMemcpy(tt->d_q,qh,(size_t)n,cudaMemcpyHostToDevice));
            ternary_optim_free(ho);
            TGK(cudaMalloc(&tt->d_scale,(size_t)tt->cols*sizeof(float)));
            tg_fill<<<tg_nb(tt->cols),256>>>(tt->d_scale,1.0f,tt->cols);
            TGK(cudaMalloc(&tt->d_ms,(size_t)tt->cols*sizeof(float))); TGK(cudaMemset(tt->d_ms,0,(size_t)tt->cols*sizeof(float)));
            TGK(cudaMalloc(&tt->d_vs,(size_t)tt->cols*sizeof(float))); TGK(cudaMemset(tt->d_vs,0,(size_t)tt->cols*sizeof(float)));
            TGK(cudaMalloc(&tt->d_tern,(size_t)n*sizeof(float)));
            TGK(cudaMalloc(&tt->d_gq,(size_t)n*sizeof(float)));
            TGK(cudaMalloc(&tt->d_sgrad,(size_t)tt->cols*sizeof(float)));
            TGK(cudaMalloc(&tt->d_s,sizeof(float)));
            tt->snap=(float*)malloc((size_t)n*sizeof(float));
            tt->hq=(int8_t*)malloc((size_t)n);
        } else {
            cont_p*p=&cont[n_cont++]; p->data=wdev; p->grad=gdev; p->n=n;
            TGK(cudaMalloc(&p->mm,(size_t)n*sizeof(float))); TGK(cudaMemset(p->mm,0,(size_t)n*sizeof(float)));
            TGK(cudaMalloc(&p->vv,(size_t)n*sizeof(float))); TGK(cudaMemset(p->vv,0,(size_t)n*sizeof(float)));
        }
    }

    /* B>1 gradient-accumulation shadow (one device buffer per registered grad). */
    float**accg=NULL;
    if(c.batch>1){ accg=(float**)calloc(m->ng,sizeof(float*));
        for(int i=0;i<m->ng;i++){ TGK(cudaMalloc(&accg[i],(size_t)m->gn[i]*sizeof(float))); } }

    /* host scratch */
    int V=m->V;
    int *in =(int*)malloc((size_t)c.batch*c.seq_len*sizeof(int));
    int *tg =(int*)malloc((size_t)c.batch*c.seq_len*sizeof(int));
    int *in1=(int*)malloc((size_t)c.seq_len*sizeof(int));
    int *tg1=(int*)malloc((size_t)c.seq_len*sizeof(int));
    int maxgn=0; for(int i=0;i<m->ng;i++) if(m->gn[i]>maxgn)maxgn=m->gn[i];
    float*hgrad=(float*)malloc((size_t)maxgn*sizeof(float));
    float*hlog =(float*)malloc((size_t)c.seq_len*V*sizeof(float));
    double*hloss=(double*)malloc((size_t)c.seq_len*sizeof(double));  /* opt#3 loss readback */

    /* ---- banner (mirror train.c) ---- */
    char isobuf[32]; iso_now(isobuf,sizeof isobuf);
    const double rand_bpb=log((double)FD_BYTE_VOCAB)/LN2;
    LOGF("# start=%s\n",isobuf);
    LOGF("# device=%s cc %d.%d  (GPU driver)\n",pr.name,pr.major,pr.minor);
    LOGF("=== RWKV-x 1M smoke on enwik8 (%s arm) [GPU] ===\n",c.use_sr4?"SR4":"AdamW");
    LOGF("config: C=%d heads=%d hs=%d ffn=%d lora=%d  layers: %d RWKV + %d attn (interleave)\n",
         c.C,c.n_head,c.head_size,c.ffn,c.lora,c.n_rwkv,c.n_attn);
    LOGF("vocab=%d seq_len=%d batch=%d steps=%d eval_every=%d eval_batches=%d seed=%llu\n",
         FD_BYTE_VOCAB,c.seq_len,c.batch,c.steps,c.eval_every,c.eval_batches,(unsigned long long)c.seed);
    LOGF("PARAM COUNT: total=%ld  (ternary=%ld, continuous=%ld)\n",
         total_params,ternary_params,total_params-ternary_params);
    if(c.use_sr4) LOGF("SR4 target_step=%.4g scale_lr=%.4g cont_lr=%.4g->%.4g  (Q_CLAMP=%d THETA=%d)\n",
                       c.target_step,c.scale_lr,c.lr0,c.lr_min,SR4_Q_CLAMP,SR4_THETA);
    else          LOGF("AdamW lr=%.4g->%.4g (cosine)\n",c.lr0,c.lr_min);
    LOGF("random-init reference bpb = ln(256)/ln2 = %.4f\n",rand_bpb);
    if(c.use_sr4) LOGF("# SR4: %d ternary matrices, %d continuous params\n",n_tern,n_cont);
    LOGF("step,train_bpb,val_bpb,val_acc,grad_norm,lr,"
         "step_sec,tok_per_sec,elapsed_sec,ETA_sec,"
         "churn,pct_neg,pct_zero,pct_pos,rail_pct,dead_pct,"
         "scale_mean,scale_min,scale_max\n");

    double best_vb=1e30; int best_vb_step=-1;
    double best_va=-1.0; int best_va_step=-1;
    int snap_valid=0;

    double t0=now_sec(); double last_log_wall=t0; int last_log_step=0;
    uint64_t counter=0;
    const float B1=0.9f,B2=0.999f,EPS=1e-8f,WD=0.0f;

    /* ================= opt#3: CUDA-graph capture of the B==1 train step =============
     * The trainer fires hundreds of tiny kernels per step; at 1M params the CPU
     * cannot issue launches fast enough and the K80 starves between them (~1% of
     * peak, launch-bound).  Capture the fixed per-step kernel sequence ONCE and
     * replay it as a single cudaGraphLaunch -> the kernels run back-to-back with
     * ~no per-launch CPU overhead.  Opt#2's resident arena already gives the stable
     * buffer pointers graphs require.  Only B==1 is graphed (the tiny-run case);
     * B>1 keeps the proven non-graph accumulation path below. */
    float *d_opt=NULL; unsigned long long *d_ctr=NULL;
    cudaGraph_t g_graph=NULL; cudaGraphExec_t g_exec=NULL; int graph_ready=0;
    TGK(cudaMalloc(&d_opt,4*sizeof(float)));            /* {lr, bc1, bc2, scale_lr/lr} */
    TGK(cudaMalloc(&d_ctr,sizeof(unsigned long long))); /* SR4 batch_counter */
    if(c.batch==1){
        /* Every kernel launches on stream 0, which -> the per-thread default stream
         * (this TU is built with --default-stream per-thread).  Route cuBLAS and
         * all per-step memcpy/memset onto that SAME per-thread stream so the entire
         * step lives on one capturable stream. */
        g_fm_stream = cudaStreamPerThread;
        FBK(cublasSetStream(m->H, cudaStreamPerThread));

        /* stage a real (step-0) batch into the resident input buffers */
        fd_byte_get_batch(&train_s,c.seed,0,0ULL,0,1,1,in,tg);
        TGK(cudaMemcpy(m->tok,in,(size_t)c.seq_len*sizeof(int),cudaMemcpyHostToDevice));
        TGK(cudaMemcpy(m->dtg,tg,(size_t)c.seq_len*sizeof(int),cudaMemcpyHostToDevice));

        /* WARMUP: one fwd+bwd actually EXECUTED (g_fm_capturing==0) to force cuBLAS
         * lazy init/workspace + arena hi-water BEFORE capture (a cudaMalloc during
         * capture would abort it).  Touches only activations/grads (recomputed every
         * step), never optimizer state, so step 0 below is still pristine. */
        fm_forward(m,in,c.seq_len);
        fm_backward(m,tg,c.seq_len,-1);
        TGK(cudaDeviceSynchronize());
        fmA_reset();   /* drop any overflow allocs + rewind the arena before capture */

        /* CAPTURE the full step: materialize + fwd + loss + bwd + optimizer.  The
         * per-step-varying scalars (cosine lr, bias-correction bc1/bc2, scale_lr/lr,
         * and the SR4 batch_counter) are read by the _g kernel variants from the
         * d_opt/d_ctr device buffers, which the driver refreshes each step before
         * replay -> fresh values every replay while pointers stay fixed. */
        g_fm_capturing = 1;
        TGK(cudaStreamBeginCapture(g_fm_stream, cudaStreamCaptureModeThreadLocal));
        if(c.use_sr4){
            for(int t=0;t<n_tern;t++){ tern_mat*tt=&tmat[t];
                tg_tern_from_q<<<tg_nb(tt->numel),256>>>(tt->d_q,tt->d_tern,SR4_THETA,tt->numel);
                tg_apply_scale<<<tg_nb((size_t)tt->numel),256>>>(tt->Weff,tt->d_tern,tt->d_scale,tt->rows,tt->cols); }
        }
        fm_forward(m,in,c.seq_len);
        fm_backward(m,tg,c.seq_len,-1);
        if(c.use_sr4){
            for(int t=0;t<n_tern;t++){ tern_mat*tt=&tmat[t];
                tg_scale_grad<<<tg_nb(tt->cols),256>>>(tt->grad,tt->d_tern,tt->d_sgrad,tt->rows,tt->cols);
                tg_gq<<<tg_nb((size_t)tt->numel),256>>>(tt->grad,tt->d_scale,tt->d_gq,tt->rows,tt->cols); }
            for(int t=0;t<n_tern;t++){ tern_mat*tt=&tmat[t];
                sr4_dev_g(tt->d_q,tt->d_gq,tt->numel,d_ctr,c.target_step,42ULL,tt->idx_base,SR4_Q_CLAMP,1,tt->d_s); }
            for(int t=0;t<n_tern;t++){ tern_mat*tt=&tmat[t];
                tg_scale_vec_g<<<tg_nb(tt->cols),256>>>(tt->d_sgrad,d_opt,tt->cols);
                adamw_dev_g(tt->d_scale,tt->d_sgrad,tt->d_ms,tt->d_vs,tt->cols,d_opt,B1,B2,EPS,WD,1); }
            for(int i=0;i<n_cont;i++){ cont_p*p=&cont[i];
                adamw_dev_g(p->data,p->grad,p->mm,p->vv,p->n,d_opt,B1,B2,EPS,WD,1); }
        } else {
            for(int i=0;i<n_cont;i++){ cont_p*p=&cont[i];
                adamw_dev_g(p->data,p->grad,p->mm,p->vv,p->n,d_opt,B1,B2,EPS,WD,1); }
        }
        cudaError_t cap_err=cudaStreamEndCapture(g_fm_stream,&g_graph);
        g_fm_capturing=0;
        if(cap_err!=cudaSuccess){ fprintf(stderr,"opt#3 graph capture failed: %s\n",cudaGetErrorString(cap_err)); return 1; }
        TGK(cudaGraphInstantiate(&g_exec,g_graph,NULL,NULL,0));
        graph_ready=1;
        LOGF("# opt#3: CUDA graph captured (full B=1 step on the per-thread stream); replaying per step.\n");
    }

    for(int step=0; step<c.steps; step++){
        float lr=cosine_lr(c.lr0,c.lr_min,step,c.steps);
        int last=(step==c.steps-1);
        int do_log=(step%c.eval_every==0||last);
        double train_ce=0.0, grad_norm=0.0;

        if(graph_ready){
            /* ---- opt#3 CUDA-graph path (B==1): one launch replays the whole step ---- */
            fd_byte_get_batch(&train_s,c.seed,0,(uint64_t)step,0,1,1,in,tg);
            TGK(cudaMemcpy(m->tok,in,(size_t)c.seq_len*sizeof(int),cudaMemcpyHostToDevice));
            TGK(cudaMemcpy(m->dtg,tg,(size_t)c.seq_len*sizeof(int),cudaMemcpyHostToDevice));
            /* refresh per-step scalars (bit-identical to the by-value non-graph path) */
            float bc1=1.0f-powf(B1,(float)(step+1)), bc2=1.0f-powf(B2,(float)(step+1));
            float h_opt[4]={ lr, bc1, bc2, (lr!=0.0f)?(c.scale_lr/lr):0.0f };
            TGK(cudaMemcpy(d_opt,h_opt,4*sizeof(float),cudaMemcpyHostToDevice));
            unsigned long long hc=(unsigned long long)step;   /* SR4 counter == step */
            TGK(cudaMemcpy(d_ctr,&hc,sizeof(hc),cudaMemcpyHostToDevice));
            TGK(cudaGraphLaunch(g_exec,g_fm_stream));
            TGK(cudaStreamSynchronize(g_fm_stream));
            /* reduce resident per-token loss -> train_ce (nv==T since ignore_index=-1) */
            TGK(cudaMemcpy(hloss,m->loss_t,(size_t)c.seq_len*sizeof(double),cudaMemcpyDeviceToHost));
            double L=0.0; for(int t=0;t<c.seq_len;t++) L+=hloss[t]; train_ce=L/(double)c.seq_len;
            /* grads survive the optimizer step unmodified, so grad_l2 post-replay matches */
            if(do_log) grad_norm=grad_l2_gpu(m,hgrad);
        } else {
            /* ---- non-graph path (B>1 accumulation, or fallback) ---- */
            if(c.use_sr4){
                for(int t=0;t<n_tern;t++){ tern_mat*tt=&tmat[t];
                    tg_tern_from_q<<<tg_nb(tt->numel),256>>>(tt->d_q,tt->d_tern,SR4_THETA,tt->numel);
                    tg_apply_scale<<<tg_nb((size_t)tt->numel),256>>>(tt->Weff,tt->d_tern,tt->d_scale,tt->rows,tt->cols);
                }
            }
            fd_byte_get_batch(&train_s,c.seed,0,(uint64_t)step,0,1,(uint32_t)c.batch,in,tg);
            if(c.batch==1){
                fm_forward(m,in,c.seq_len);
                train_ce=fm_backward(m,tg,c.seq_len,-1);
            } else {
                for(int i=0;i<m->ng;i++) TGK(cudaMemset(accg[i],0,(size_t)m->gn[i]*sizeof(float)));
                for(int b=0;b<c.batch;b++){
                    memcpy(in1,&in[(size_t)b*c.seq_len],(size_t)c.seq_len*sizeof(int));
                    memcpy(tg1,&tg[(size_t)b*c.seq_len],(size_t)c.seq_len*sizeof(int));
                    fm_forward(m,in1,c.seq_len);
                    train_ce+=fm_backward(m,tg1,c.seq_len,-1);
                    for(int i=0;i<m->ng;i++) tg_acc_add<<<tg_nb((size_t)m->gn[i]),256>>>(accg[i],m->gptr[i],(size_t)m->gn[i]);
                }
                train_ce/=c.batch;
                const float inv=1.0f/(float)c.batch;
                for(int i=0;i<m->ng;i++) tg_scale_write<<<tg_nb((size_t)m->gn[i]),256>>>(m->gptr[i],accg[i],inv,(size_t)m->gn[i]);
            }
            grad_norm=do_log?grad_l2_gpu(m,hgrad):0.0;
            if(c.use_sr4){
                for(int t=0;t<n_tern;t++){ tern_mat*tt=&tmat[t];
                    tg_scale_grad<<<tg_nb(tt->cols),256>>>(tt->grad,tt->d_tern,tt->d_sgrad,tt->rows,tt->cols);
                    tg_gq<<<tg_nb((size_t)tt->numel),256>>>(tt->grad,tt->d_scale,tt->d_gq,tt->rows,tt->cols);
                }
                for(int t=0;t<n_tern;t++){ tern_mat*tt=&tmat[t];
                    sr4_dev(tt->d_q,tt->d_gq,tt->numel,counter,c.target_step,42ULL,tt->idx_base,SR4_Q_CLAMP,1,tt->d_s); }
                counter++;
                for(int t=0;t<n_tern;t++){ tern_mat*tt=&tmat[t];
                    if(c.scale_lr!=lr) tg_scale_vec<<<tg_nb(tt->cols),256>>>(tt->d_sgrad,c.scale_lr/lr,tt->cols);
                    adamw_dev(tt->d_scale,tt->d_sgrad,tt->d_ms,tt->d_vs,tt->cols,step+1,lr,B1,B2,EPS,WD,1); }
                for(int i=0;i<n_cont;i++){ cont_p*p=&cont[i];
                    adamw_dev(p->data,p->grad,p->mm,p->vv,p->n,step+1,lr,B1,B2,EPS,WD,1); }
            } else {
                for(int i=0;i<n_cont;i++){ cont_p*p=&cont[i];
                    adamw_dev(p->data,p->grad,p->mm,p->vv,p->n,step+1,lr,B1,B2,EPS,WD,1); }
            }
            TGK(cudaDeviceSynchronize());
        }

        if(do_log){
            double vacc=0.0;
            double vb=eval_metrics_gpu(m,&val_s,c.seed,c.seq_len,c.eval_batches,in1,tg1,hlog,&vacc);
            double tb=train_ce/LN2;
            double now=now_sec(); double elapsed=now-t0;
            int dstep=step-last_log_step;
            double step_sec=(dstep>0)?(now-last_log_wall)/(double)dstep:elapsed/(double)(step+1);
            double avg_step=elapsed/(double)(step+1);
            double tok_per_sec=step_sec>0.0?(double)c.batch*(double)c.seq_len/step_sec:0.0;
            double eta=(double)(c.steps-1-step)*avg_step;
            last_log_wall=now; last_log_step=step;
            if(vb<best_vb){best_vb=vb;best_vb_step=step;}
            if(vacc>best_va){best_va=vacc;best_va_step=step;}
            if(c.use_sr4){
                sr4_panel p; sr4_panel_compute(tmat,n_tern,snap_valid,&p,NULL); snap_valid=1;
                LOGF("%d,%.5f,%.5f,%.5f,%.4g,%.4g,"
                     "%.4f,%.1f,%.2f,%.2f,"
                     "%.5f,%.5f,%.5f,%.5f,%.5f,%.5f,"
                     "%.4f,%.4f,%.4f\n",
                     step,tb,vb,vacc,grad_norm,lr,step_sec,tok_per_sec,elapsed,eta,
                     p.churn,p.pct_neg,p.pct_zero,p.pct_pos,p.rail_pct,p.dead_pct,
                     p.scale_mean,p.scale_min,p.scale_max);
            } else {
                LOGF("%d,%.5f,%.5f,%.5f,%.4g,%.4g,"
                     "%.4f,%.1f,%.2f,%.2f,"
                     "-,-,-,-,-,-,-,-,-\n",
                     step,tb,vb,vacc,grad_norm,lr,step_sec,tok_per_sec,elapsed,eta);
            }
            LOGF("# BEST: val_bpb=%.5f (step %d) val_acc=%.5f (step %d)\n",best_vb,best_vb_step,best_va,best_va_step);
        }
    }
    double wall=now_sec()-t0;

    /* final metrics (mirror train.c) */
    double final_train_bpb,final_val_bpb,final_val_acc=0.0;
    if(c.use_sr4){
        for(int t=0;t<n_tern;t++){ tern_mat*tt=&tmat[t];
            tg_tern_from_q<<<tg_nb(tt->numel),256>>>(tt->d_q,tt->d_tern,SR4_THETA,tt->numel);
            tg_apply_scale<<<tg_nb((size_t)tt->numel),256>>>(tt->Weff,tt->d_tern,tt->d_scale,tt->rows,tt->cols); }
    }
    fd_byte_get_batch(&train_s,c.seed,0,(uint64_t)c.steps,0,1,1,in1,tg1);
    fm_forward(m,in1,c.seq_len);
    TGK(cudaMemcpy(hlog,m->logits,(size_t)c.seq_len*V*sizeof(float),cudaMemcpyDeviceToHost));
    final_train_bpb=rx_cross_entropy_forward(hlog,tg1,c.seq_len,V,-1)/LN2;
    final_val_bpb=eval_metrics_gpu(m,&val_s,c.seed,c.seq_len,c.eval_batches*2,in1,tg1,hlog,&final_val_acc);

    LOGF("# FINAL: train_bpb=%.5f val_bpb=%.5f val_acc=%.5f  wall=%.2fs  (%.3fs/step over %d steps)\n",
         final_train_bpb,final_val_bpb,final_val_acc,wall,wall/c.steps,c.steps);
    LOGF("# BEST: val_bpb=%.5f (step %d) val_acc=%.5f (step %d)\n",best_vb,best_vb_step,best_va,best_va_step);
    {
        double avg_step_sec=c.steps>0?wall/(double)c.steps:0.0;
        long total_tokens=(long)c.steps*(long)c.batch*(long)c.seq_len;
        double tps=wall>0.0?(double)total_tokens/wall:0.0;
        char endiso[32]; iso_now(endiso,sizeof endiso);
        LOGF("# ===== END-OF-RUN SUMMARY (extrapolate future-run wall-time from these) =====\n");
        LOGF("# end=%s\n",endiso);
        LOGF("# total_wall_sec=%.2f  total_steps=%d  avg_step_sec=%.4f  tok_per_sec=%.1f  total_tokens=%ld  (batch=%d seq_len=%d)\n",
             wall,c.steps,avg_step_sec,tps,total_tokens,c.batch,c.seq_len);
    }

    if(logfp) fclose(logfp);
    return 0;
}