/* SPDX-License-Identifier: AGPL-3.0-or-later */
/* Copyright (C) 2026 David (Corey) Herrera */
/*
 * timemix.cu — RWKV-7 time-mix block forward+backward on the K80 (sm_37).
 *
 * GPU counterpart of the CPU oracle rx_timemix_forward_seq / rx_timemix_backward
 * (rwkvx/rx_pretrain.c).  ALL matmuls (r/k/v/output projections + w/a/g/v LoRA
 * pairs) go through cuBLAS SGEMM over the whole [B*T, *] token batch; every
 * non-matmul op (token-shift/lerp, decay exp, sigmoid/tanh gates, L2-norm per
 * head, key-mod, value-residual, WKV7 composites, group-norm, r*k*r_k bonus,
 * g-gate, output blend) is an elementwise CUDA kernel; the recurrence in the
 * middle is the N=64-safe wkv7_fwd_kernel / wkv7_bwd_kernel.
 *
 * cuBLAS is column-major; foundry weights are row-major [in,out].  Row-major
 * C[M,N] = op(A)[M,K] @ op(B)[K,N] is issued as the column-major transpose
 * C^T = op(B)^T @ op(A)^T, i.e. cublasSgemm(opB,opA, N,M,K, B,ldb, A,lda, C,N)
 * with A,B SWAPPED and ldb=(tB?K:N), lda=(tA?M:K), ldc=N (see rm_gemm).
 *
 * Layouts identical to the oracle: token vectors [B][T][C] row-major
 * (C=H*N); WKV7 vectors [B][T][H][N]; state [B][H][N][N].  Since C = H*N and
 * per-token the head-major channel order matches (i -> h*N+n), the [B][T][C]
 * token buffers ARE the [B][T][H][N] WKV7 buffers with no reshape.
 */
#include <cuda_runtime.h>
#include <cublas_v2.h>
#include <stdio.h>
#include <stdlib.h>
#include <math.h>

#include "cuda_kernels/wkv7.cu"   /* wkv7_fwd_kernel / wkv7_bwd_kernel + launchers */
#include "cuda_kernels/fm_arena.cuh" /* device scratch bump-arena (opt #2) */
#include "rwkvx/rx_pretrain.h"    /* rx_tm_params / rx_tm_grads */

#define TCK(call) do{ cudaError_t e_=(call); if(e_!=cudaSuccess){ \
  fprintf(stderr,"CUDA %s:%d %s\n",__FILE__,__LINE__,cudaGetErrorString(e_)); return -1; } }while(0)
#define TBK(call) do{ cublasStatus_t s_=(call); if(s_!=CUBLAS_STATUS_SUCCESS){ \
  fprintf(stderr,"cuBLAS %s:%d status %d\n",__FILE__,__LINE__,(int)s_); return -1; } }while(0)

#define RWDECAY_COEFF 0.606531f
#define TPB 256
static inline int nblk(size_t n){ return (int)((n+TPB-1)/TPB); }

/* ---- row-major GEMM via cuBLAS: C[M,N] = op(A) @ op(B) ---- */
static cublasStatus_t rm_gemm(cublasHandle_t H,int tA,int tB,int M,int N,int K,
    float alpha,const float*A,const float*B,float beta,float*C){
    int lda = tA? M : K;
    int ldb = tB? K : N;
    return cublasSgemm(H, tB?CUBLAS_OP_T:CUBLAS_OP_N, tA?CUBLAS_OP_T:CUBLAS_OP_N,
                       N, M, K, &alpha, B, ldb, A, lda, &beta, C, N);
}

/* =============================== forward kernels =============================== */
__global__ void k_tokenshift(const float*x,const float*xprev,
    const float*xr,const float*xk,const float*xv,const float*xw,const float*xa,const float*xg,
    float*shr,float*shk,float*shv,float*shw,float*sha,float*shg,int B,int T,int C){
    size_t g=(size_t)blockIdx.x*blockDim.x+threadIdx.x; size_t tot=(size_t)B*T*C; if(g>=tot)return;
    int i=g%C; size_t token=g/C; int t=(int)(token%T); size_t b=token/T;
    float xx=x[g];
    float prev = (t==0)? (xprev? xprev[b*C+i] : 0.0f) : x[g-C];
    float d=prev-xx;
    shr[g]=xx+d*xr[i]; shk[g]=xx+d*xk[i]; shv[g]=xx+d*xv[i];
    shw[g]=xx+d*xw[i]; sha[g]=xx+d*xa[i]; shg[g]=xx+d*xg[i];
}
__global__ void k_decay(const float*w0,const float*ow,float*sigw,float*W,int C,size_t tot){
    size_t g=(size_t)blockIdx.x*blockDim.x+threadIdx.x; if(g>=tot)return;
    float sig=1.0f/(1.0f+expf(-(w0[g%C]+ow[g]))); sigw[g]=sig; W[g]=expf(-RWDECAY_COEFF*sig);
}
__global__ void k_bias_sigmoid(const float*bias,float*io,int C,size_t tot){
    size_t g=(size_t)blockIdx.x*blockDim.x+threadIdx.x; if(g>=tot)return;
    io[g]=1.0f/(1.0f+expf(-(bias[g%C]+io[g])));
}
__global__ void k_tanh(float*io,size_t tot){
    size_t g=(size_t)blockIdx.x*blockDim.x+threadIdx.x; if(g>=tot)return; io[g]=tanhf(io[g]); }
__global__ void k_sigmoid(float*io,size_t tot){
    size_t g=(size_t)blockIdx.x*blockDim.x+threadIdx.x; if(g>=tot)return; io[g]=1.0f/(1.0f+expf(-io[g])); }
__global__ void k_tk(const float*kvec,const float*k_k,float*tk,int C,size_t tot){
    size_t g=(size_t)blockIdx.x*blockDim.x+threadIdx.x; if(g>=tot)return; tk[g]=kvec[g]*k_k[g%C]; }
__global__ void k_l2norm(const float*tk,float*kk,int H,int N,size_t BT){
    size_t th=(size_t)blockIdx.x*blockDim.x+threadIdx.x; if(th>=BT*H)return;
    const float*t=tk+th*N; float*o=kk+th*N; float ss=0.0f;
    for(int j=0;j<N;j++) ss+=t[j]*t[j];
    float inv=1.0f/(sqrtf(ss)+1e-12f);
    for(int j=0;j<N;j++) o[j]=t[j]*inv;
}
__global__ void k_keymod(const float*kvec,const float*avec,const float*k_a,float*K,int C,size_t tot){
    size_t g=(size_t)blockIdx.x*blockDim.x+threadIdx.x; if(g>=tot)return;
    int i=g%C; K[g]=kvec[g]*(1.0f+(avec[g]-1.0f)*k_a[i]);
}
__global__ void k_vres0(const float*vvec,float*V,float*vfirst,size_t tot){
    size_t g=(size_t)blockIdx.x*blockDim.x+threadIdx.x; if(g>=tot)return; V[g]=vvec[g]; vfirst[g]=vvec[g]; }
__global__ void k_vres(const float*vvec,const float*vfirst,const float*v0,float*ov_blend,float*V,int C,size_t tot){
    size_t g=(size_t)blockIdx.x*blockDim.x+threadIdx.x; if(g>=tot)return;
    float bl=1.0f/(1.0f+expf(-(v0[g%C]+ov_blend[g]))); ov_blend[g]=bl;
    V[g]=vvec[g]+(vfirst[g]-vvec[g])*bl;
}
__global__ void k_composites(const float*kk,const float*avec,float*A,float*Bb,size_t tot){
    size_t g=(size_t)blockIdx.x*blockDim.x+threadIdx.x; if(g>=tot)return; A[g]=-kk[g]; Bb[g]=kk[g]*avec[g]; }
__global__ void k_groupnorm(const float*y,const float*w,const float*b,float*ygn,
    float*mean_o,float*rstd_o,int H,int N,size_t BT){
    size_t th=(size_t)blockIdx.x*blockDim.x+threadIdx.x; if(th>=BT*H)return;
    int h=(int)(th%H); const float*gy=y+th*N; float*go=ygn+th*N;
    float mean=0.0f; for(int j=0;j<N;j++) mean+=gy[j]; mean/=N;
    float var=0.0f; for(int j=0;j<N;j++){float d=gy[j]-mean; var+=d*d;} var/=N;
    float rstd=1.0f/sqrtf(var+64e-5f); mean_o[th]=mean; rstd_o[th]=rstd;
    int off=h*N; for(int j=0;j<N;j++) go[j]=(gy[j]-mean)*rstd*w[off+j]+b[off+j];
}
__global__ void k_bonus(const float*R,const float*K,const float*r_k,const float*V,
    const float*ygn,float*yb,float*bonus,int H,int N,size_t BT){
    size_t th=(size_t)blockIdx.x*blockDim.x+threadIdx.x; if(th>=BT*H)return;
    int h=(int)(th%H); int off=h*N; size_t base=th*N;
    float s=0.0f; for(int j=0;j<N;j++) s+=R[base+j]*K[base+j]*r_k[off+j];
    bonus[th]=s; for(int j=0;j<N;j++) yb[base+j]=ygn[base+j]+s*V[base+j];
}
__global__ void k_mul(const float*a,const float*b,float*o,size_t tot){
    size_t g=(size_t)blockIdx.x*blockDim.x+threadIdx.x; if(g>=tot)return; o[g]=a[g]*b[g]; }

/* =============================== device weight/cache holders =============================== */
typedef struct {
    float *x_r,*x_k,*x_v,*x_w,*x_a,*x_g;
    float *receptance_w,*key_w,*value_w,*output_w;
    float *w0,*w1,*w2,*a0,*a1,*a2,*g1,*g2,*v0,*v1,*v2;
    float *k_k,*k_a,*r_k,*ln_x_w,*ln_x_b;
} tm_dw;

static float* up(const float*h,size_t n){ /* upload host->device (NULL passthrough) */
    if(!h) return NULL; float*d=NULL; if(cudaMalloc(&d,n*sizeof(float))!=cudaSuccess) return NULL;
    if(cudaMemcpy(d,h,n*sizeof(float),cudaMemcpyHostToDevice)!=cudaSuccess){cudaFree(d);return NULL;} return d; }

static int tm_upload_weights(const rx_tm_params*P, tm_dw*D){
    int C=P->c,Dw=P->Dw,Da=P->Da,Dg=P->Dg,Dv=P->Dv;
    D->x_r=up(P->x_r,C);D->x_k=up(P->x_k,C);D->x_v=up(P->x_v,C);
    D->x_w=up(P->x_w,C);D->x_a=up(P->x_a,C);D->x_g=up(P->x_g,C);
    D->receptance_w=up(P->receptance_w,(size_t)C*C);D->key_w=up(P->key_w,(size_t)C*C);
    D->value_w=up(P->value_w,(size_t)C*C);D->output_w=up(P->output_w,(size_t)C*C);
    D->w0=up(P->w0,C);D->w1=up(P->w1,(size_t)C*Dw);D->w2=up(P->w2,(size_t)Dw*C);
    D->a0=up(P->a0,C);D->a1=up(P->a1,(size_t)C*Da);D->a2=up(P->a2,(size_t)Da*C);
    D->g1=up(P->g1,(size_t)C*Dg);D->g2=up(P->g2,(size_t)Dg*C);
    if(P->layer_id>0){D->v0=up(P->v0,C);D->v1=up(P->v1,(size_t)C*Dv);D->v2=up(P->v2,(size_t)Dv*C);}
    else {D->v0=D->v1=D->v2=NULL;}
    D->k_k=up(P->k_k,C);D->k_a=up(P->k_a,C);D->r_k=up(P->r_k,C);
    D->ln_x_w=up(P->ln_x_w,C);D->ln_x_b=up(P->ln_x_b,C);
    return 0;
}
static void tm_free_weights(tm_dw*D){
    float**p=(float**)D; for(unsigned i=0;i<sizeof(tm_dw)/sizeof(float*);i++) if(p[i]) cudaFree(p[i]);
}

/* forward activation caches (device), all sized against BT/C/lora */
typedef struct {
    float *shr,*shk,*shv,*shw,*sha,*shg;
    float *R,*kvec,*vvec,*sigw,*avec,*gvec,*tk,*kk,*K,*V,*A,*Bb,*W;
    float *lw,*tw,*la,*oa,*lg,*sg,*lv,*ov,*blend;
    float *y,*ygn,*yb,*tg,*gn_mean,*gn_rstd,*bonus;
} tm_cache_d;

static void tmc_free(tm_cache_d*c){
    float**p=(float**)c; for(unsigned i=0;i<sizeof(tm_cache_d)/sizeof(float*);i++) if(p[i]) fmA_free(p[i]);
}

/* run the forward on device weights D + device x(dx_in), fill caches C_ and dv_first, produce d_out */
static int tm_fwd_device(const rx_tm_params*P, cublasHandle_t H, const tm_dw*D,
    const float*x, const float*xprev, float*vfirst, tm_cache_d*c, float*out, int B,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;
    size_t BT=(size_t)B*T, tot=BT*C;
    #define AL(f,n) do{ c->f=(float*)fmA_alloc((size_t)(n)*sizeof(float)); }while(0)
    AL(shr,tot);AL(shk,tot);AL(shv,tot);AL(shw,tot);AL(sha,tot);AL(shg,tot);
    AL(R,tot);AL(kvec,tot);AL(vvec,tot);AL(sigw,tot);AL(avec,tot);AL(gvec,tot);
    AL(tk,tot);AL(kk,tot);AL(K,tot);AL(V,tot);AL(A,tot);AL(Bb,tot);AL(W,tot);
    AL(lw,BT*Dw);AL(tw,BT*Dw);AL(la,BT*Da);AL(oa,tot);AL(lg,BT*Dg);AL(sg,BT*Dg);
    AL(lv,BT*(Dv>0?Dv:1));AL(ov,tot);AL(blend,tot);
    AL(y,tot);AL(ygn,tot);AL(yb,tot);AL(tg,tot);AL(gn_mean,BT*nH);AL(gn_rstd,BT*nH);AL(bonus,BT*nH);
    #undef AL

    k_tokenshift<<<nblk(tot),TPB>>>(x,xprev,D->x_r,D->x_k,D->x_v,D->x_w,D->x_a,D->x_g,
        c->shr,c->shk,c->shv,c->shw,c->sha,c->shg,B,T,C);
    /* r,k,v projections */
    TBK(rm_gemm(H,0,0,BT,C,C,1.0f,c->shr,D->receptance_w,0.0f,c->R));
    TBK(rm_gemm(H,0,0,BT,C,C,1.0f,c->shk,D->key_w,0.0f,c->kvec));
    TBK(rm_gemm(H,0,0,BT,C,C,1.0f,c->shv,D->value_w,0.0f,c->vvec));
    /* decay w */
    TBK(rm_gemm(H,0,0,BT,Dw,C,1.0f,c->shw,D->w1,0.0f,c->lw));
    TCK(cudaMemcpyAsync(c->tw,c->lw,BT*Dw*sizeof(float),cudaMemcpyDeviceToDevice,g_fm_stream)); /* opt#3: capture-safe */
    k_tanh<<<nblk(BT*Dw),TPB>>>(c->tw,BT*Dw);
    { float*ow=(float*)fmA_alloc(tot*sizeof(float));
      TBK(rm_gemm(H,0,0,BT,C,Dw,1.0f,c->tw,D->w2,0.0f,ow));
      k_decay<<<nblk(tot),TPB>>>(D->w0,ow,c->sigw,c->W,C,tot); fmA_free(ow); }
    /* gate a */
    TBK(rm_gemm(H,0,0,BT,Da,C,1.0f,c->sha,D->a1,0.0f,c->la));
    TBK(rm_gemm(H,0,0,BT,C,Da,1.0f,c->la,D->a2,0.0f,c->avec));
    k_bias_sigmoid<<<nblk(tot),TPB>>>(D->a0,c->avec,C,tot);
    /* gate g */
    TBK(rm_gemm(H,0,0,BT,Dg,C,1.0f,c->shg,D->g1,0.0f,c->lg));
    TCK(cudaMemcpyAsync(c->sg,c->lg,BT*Dg*sizeof(float),cudaMemcpyDeviceToDevice,g_fm_stream)); /* opt#3: capture-safe */
    k_sigmoid<<<nblk(BT*Dg),TPB>>>(c->sg,BT*Dg);
    TBK(rm_gemm(H,0,0,BT,C,Dg,1.0f,c->sg,D->g2,0.0f,c->gvec));
    /* key norm + key mod */
    k_tk<<<nblk(tot),TPB>>>(c->kvec,D->k_k,c->tk,C,tot);
    k_l2norm<<<nblk(BT*nH),TPB>>>(c->tk,c->kk,nH,N,BT);
    k_keymod<<<nblk(tot),TPB>>>(c->kvec,c->avec,D->k_a,c->K,C,tot);
    /* R output = rvec (already in c->R) */
    /* value residual */
    if(P->layer_id==0){ k_vres0<<<nblk(tot),TPB>>>(c->vvec,c->V,vfirst,tot); }
    else {
        TBK(rm_gemm(H,0,0,BT,Dv,C,1.0f,c->shv,D->v1,0.0f,c->lv));
        TBK(rm_gemm(H,0,0,BT,C,Dv,1.0f,c->lv,D->v2,0.0f,c->blend));
        k_vres<<<nblk(tot),TPB>>>(c->vvec,vfirst,D->v0,c->blend,c->V,C,tot);
    }
    /* composites + WKV7 */
    k_composites<<<nblk(tot),TPB>>>(c->kk,c->avec,c->A,c->Bb,tot);
    { size_t stsz=(size_t)B*nH*N*N, sasz=(size_t)T*B*nH*N, scsz=(size_t)T*stsz;
      float*SC=(float*)fmA_alloc(scsz*sizeof(float)); float*SA=(float*)fmA_alloc(sasz*sizeof(float));
      size_t shf=((size_t)N*N+7*N)*sizeof(float);
      wkv7_fwd_kernel<<<B*nH,N,shf>>>(c->R,c->W,c->K,c->V,c->A,c->Bb,NULL,c->y,NULL,SC,SA,B,T,nH,N);
      fmA_free(SC); fmA_free(SA); }
    /* group-norm + bonus + g-gate + output */
    k_groupnorm<<<nblk(BT*nH),TPB>>>(c->y,D->ln_x_w,D->ln_x_b,c->ygn,c->gn_mean,c->gn_rstd,nH,N,BT);
    k_bonus<<<nblk(BT*nH),TPB>>>(c->R,c->K,D->r_k,c->V,c->ygn,c->yb,c->bonus,nH,N,BT);
    k_mul<<<nblk(tot),TPB>>>(c->yb,c->gvec,c->tg,tot);
    TBK(rm_gemm(H,0,0,BT,C,C,1.0f,c->tg,D->output_w,0.0f,out));
    fm_sync();   /* opt#3: no-op during graph capture */
    return 0;
}

/* ---- host-facing forward parity entry: host in, host out ---- */
extern "C" int timemix_forward_cuda(const rx_tm_params*P,const float*x,const float*x_prev_init,
    float*v_first,float*out,int B,int T){
    int C=P->c; size_t BT=(size_t)B*T, tot=BT*C;
    cublasHandle_t H; if(cublasCreate(&H)!=CUBLAS_STATUS_SUCCESS) return -1;
    tm_dw D; if(tm_upload_weights(P,&D)) return -1;
    float*dx,*dxprev=NULL,*dvf,*dout;
    TCK(cudaMalloc(&dx,tot*sizeof(float))); TCK(cudaMemcpy(dx,x,tot*sizeof(float),cudaMemcpyHostToDevice));
    if(x_prev_init){ TCK(cudaMalloc(&dxprev,(size_t)B*C*sizeof(float)));
        TCK(cudaMemcpy(dxprev,x_prev_init,(size_t)B*C*sizeof(float),cudaMemcpyHostToDevice)); }
    TCK(cudaMalloc(&dvf,tot*sizeof(float)));
    if(P->layer_id>0) TCK(cudaMemcpy(dvf,v_first,tot*sizeof(float),cudaMemcpyHostToDevice));
    TCK(cudaMalloc(&dout,tot*sizeof(float)));
    tm_cache_d c; memset(&c,0,sizeof(c));
    int rc=tm_fwd_device(P,H,&D,dx,dxprev,dvf,&c,dout,B,T);
    if(rc==0){
        TCK(cudaMemcpy(out,dout,tot*sizeof(float),cudaMemcpyDeviceToHost));
        if(P->layer_id==0) TCK(cudaMemcpy(v_first,dvf,tot*sizeof(float),cudaMemcpyDeviceToHost));
    }
    tmc_free(&c); tm_free_weights(&D);
    cudaFree(dx); if(dxprev)cudaFree(dxprev); cudaFree(dvf); cudaFree(dout);
    cublasDestroy(H);
    return rc;
}

/* =============================== backward kernels =============================== */
__global__ void k_split(const float*dtg,const float*gvec,const float*yb,
    float*dyb,float*dg_post,size_t tot){
    size_t g=(size_t)blockIdx.x*blockDim.x+threadIdx.x; if(g>=tot)return;
    dyb[g]=dtg[g]*gvec[g]; dg_post[g]=dtg[g]*yb[g];
}
/* bonus backward: db_h=sum_j dyb*V; dvf_post,dr_post,dkmod_post; G_r_k atomic */
__global__ void k_bonus_bwd(const float*dyb,const float*V,const float*R,const float*K,
    const float*r_k,const float*bonus,float*dvf_post,float*dr_post,float*dkmod_post,
    float*G_r_k,int H,int N,size_t BT){
    size_t th=(size_t)blockIdx.x*blockDim.x+threadIdx.x; if(th>=BT*H)return;
    int h=(int)(th%H); int off=h*N; size_t base=th*N; float bon=bonus[th];
    float db=0.0f; for(int j=0;j<N;j++) db+=dyb[base+j]*V[base+j];
    for(int j=0;j<N;j++){
        float rj=R[base+j], kj=K[base+j], rk=r_k[off+j];
        dvf_post[base+j]  = dyb[base+j]*bon;
        dr_post[base+j]   = db*kj*rk;
        dkmod_post[base+j]= db*rj*rk;
        atomicAdd(&G_r_k[off+j], db*rj*kj);
    }
}
__global__ void k_groupnorm_bwd(const float*y,const float*dygn,const float*ln_x_w,
    const float*mean_o,const float*rstd_o,float*dy_wkv,float*G_ln_w,float*G_ln_b,
    int H,int N,size_t BT){
    size_t th=(size_t)blockIdx.x*blockDim.x+threadIdx.x; if(th>=BT*H)return;
    int h=(int)(th%H); int off=h*N; size_t base=th*N;
    float mean=mean_o[th], rstd=rstd_o[th];
    float mdn=0.0f, mdnn=0.0f;
    for(int j=0;j<N;j++){
        float nv=(y[base+j]-mean)*rstd;
        float dnorm=dygn[base+j]*ln_x_w[off+j];
        mdn+=dnorm; mdnn+=dnorm*nv;
        atomicAdd(&G_ln_w[off+j], dygn[base+j]*nv);
        atomicAdd(&G_ln_b[off+j], dygn[base+j]);
    }
    mdn/=N; mdnn/=N;
    for(int j=0;j<N;j++){
        float nv=(y[base+j]-mean)*rstd;
        float dnorm=dygn[base+j]*ln_x_w[off+j];
        dy_wkv[base+j]=rstd*(dnorm - mdn - nv*mdnn);
    }
}
__global__ void k_p4_seed(const float*gR,const float*dr_post,const float*gK,const float*dkmod_post,
    const float*gV,const float*dvf_post,float*dr_t,float*dkmod_t,float*dvf_t,size_t tot){
    size_t g=(size_t)blockIdx.x*blockDim.x+threadIdx.x; if(g>=tot)return;
    dr_t[g]=gR[g]+dr_post[g]; dkmod_t[g]=gK[g]+dkmod_post[g]; dvf_t[g]=gV[g]+dvf_post[g];
}
__global__ void k_comp_keymod(const float*gA,const float*gB,const float*avec,const float*kk,
    const float*kvec,const float*k_a,const float*dkmod_t,float*dkk,float*dav,float*dkvec,
    float*G_k_a,int C,size_t tot){
    size_t g=(size_t)blockIdx.x*blockDim.x+threadIdx.x; if(g>=tot)return; int i=g%C;
    dkk[g]=-gA[g]+avec[g]*gB[g];
    float factor=1.0f+(avec[g]-1.0f)*k_a[i];
    dkvec[g]=dkmod_t[g]*factor;
    dav[g]=kk[g]*gB[g]+dkmod_t[g]*kvec[g]*k_a[i];
    atomicAdd(&G_k_a[i], dkmod_t[g]*kvec[g]*(avec[g]-1.0f));
}
__global__ void k_l2norm_bwd(const float*tk,const float*dkk,float*dtk,int H,int N,size_t BT){
    size_t th=(size_t)blockIdx.x*blockDim.x+threadIdx.x; if(th>=BT*H)return;
    size_t base=th*N;
    double ss=0.0; for(int j=0;j<N;j++){double u=tk[base+j]; ss+=u*u;}
    double norm=sqrt(ss), inv=1.0/(norm+1e-12);
    double gdot=0.0; for(int j=0;j<N;j++) gdot+=(double)dkk[base+j]*(double)tk[base+j];
    for(int j=0;j<N;j++){ double u=tk[base+j]; double d=inv*(double)dkk[base+j];
        if(norm>0.0) d-=inv*inv*gdot*u/norm; dtk[base+j]=(float)d; }
}
__global__ void k_tk_bwd(const float*dtk,const float*k_k,const float*kvec,float*dkvec,
    float*G_k_k,int C,size_t tot){
    size_t g=(size_t)blockIdx.x*blockDim.x+threadIdx.x; if(g>=tot)return; int i=g%C;
    dkvec[g]+=dtk[g]*k_k[i]; atomicAdd(&G_k_k[i], dtk[g]*kvec[g]);
}
__global__ void k_avec_bwd(const float*dav,const float*avec,float*d_oa,float*G_a0,int C,size_t tot){
    size_t g=(size_t)blockIdx.x*blockDim.x+threadIdx.x; if(g>=tot)return; int i=g%C;
    float av=avec[g]; float d=dav[g]*av*(1.0f-av); d_oa[g]=d; atomicAdd(&G_a0[i], d);
}
__global__ void k_wdecay_bwd(const float*gW,const float*W,const float*sigw,float*d_ow,
    float*G_w0,int C,size_t tot){
    size_t g=(size_t)blockIdx.x*blockDim.x+threadIdx.x; if(g>=tot)return; int i=g%C;
    float dsig=gW[g]*W[g]*(-RWDECAY_COEFF); float sg=sigw[g];
    float d=dsig*sg*(1.0f-sg); d_ow[g]=d; atomicAdd(&G_w0[i], d);
}
__global__ void k_tanhderiv(const float*d_tw,const float*tw,float*d_lw,size_t n){
    size_t g=(size_t)blockIdx.x*blockDim.x+threadIdx.x; if(g>=n)return;
    float tv=tw[g]; d_lw[g]=d_tw[g]*(1.0f-tv*tv);
}
__global__ void k_sigderiv(const float*d_sg,const float*sg,float*d_lg,size_t n){
    size_t g=(size_t)blockIdx.x*blockDim.x+threadIdx.x; if(g>=n)return;
    float s=sg[g]; d_lg[g]=d_sg[g]*s*(1.0f-s);
}
__global__ void k_vvec0(const float*dvf_t,const float*dvf_in,float*dvvec,size_t tot){
    size_t g=(size_t)blockIdx.x*blockDim.x+threadIdx.x; if(g>=tot)return;
    dvvec[g]=dvf_t[g]+(dvf_in?dvf_in[g]:0.0f);
}
__global__ void k_vvec1(const float*dvf_t,const float*blend,const float*vf,const float*vvec,
    const float*v0,float*dvvec,float*d_vf_out,float*d_ov,float*G_v0,int C,size_t tot){
    size_t g=(size_t)blockIdx.x*blockDim.x+threadIdx.x; if(g>=tot)return; int i=g%C;
    float bl=blend[g]; dvvec[g]=dvf_t[g]*(1.0f-bl);
    if(d_vf_out) d_vf_out[g]=dvf_t[g]*bl;
    float d=dvf_t[g]*(vf[g]-vvec[g])*bl*(1.0f-bl); d_ov[g]=d; atomicAdd(&G_v0[i], d);
}
__global__ void k_tokenshift_bwd(const float*x,const float*xprev,
    const float*dsr,const float*dsk,const float*dsv,const float*dsw,const float*dsa,const float*dsg,
    const float*xr,const float*xk,const float*xv,const float*xw,const float*xa,const float*xg,
    float*gxr,float*gxk,float*gxv,float*gxw,float*gxa,float*gxg,
    float*dx,float*dxprev,int B,int T,int C){
    size_t g=(size_t)blockIdx.x*blockDim.x+threadIdx.x; size_t tot=(size_t)B*T*C; if(g>=tot)return;
    int i=g%C; size_t token=g/C; int t=(int)(token%T); size_t b=token/T;
    float xx=x[g];
    float prev=(t==0)?(xprev?xprev[b*C+i]:0.0f):x[g-C];
    float pmx=prev-xx;
    atomicAdd(&gxr[i], dsr[g]*pmx); atomicAdd(&gxk[i], dsk[g]*pmx); atomicAdd(&gxv[i], dsv[g]*pmx);
    atomicAdd(&gxw[i], dsw[g]*pmx); atomicAdd(&gxa[i], dsa[g]*pmx); atomicAdd(&gxg[i], dsg[g]*pmx);
    float own = dsr[g]*(1.0f-xr[i])+dsk[g]*(1.0f-xk[i])+dsv[g]*(1.0f-xv[i])
              + dsw[g]*(1.0f-xw[i])+dsa[g]*(1.0f-xa[i])+dsg[g]*(1.0f-xg[i]);
    float nxt=0.0f;
    if(t<T-1){ size_t gn=g+C;
        nxt = dsr[gn]*xr[i]+dsk[gn]*xk[i]+dsv[gn]*xv[i]+dsw[gn]*xw[i]+dsa[gn]*xa[i]+dsg[gn]*xg[i]; }
    dx[g]=own+nxt;
    if(t==0 && dxprev){
        dxprev[b*C+i]=dsr[g]*xr[i]+dsk[g]*xk[i]+dsv[g]*xv[i]+dsw[g]*xw[i]+dsa[g]*xa[i]+dsg[g]*xg[i];
    }
}

/* =============================== backward driver =============================== */
extern "C" int timemix_backward_cuda(const rx_tm_params*P,const float*x,const float*x_prev_init,
    const float*v_first,const float*dy,rx_tm_grads*G,float*dx,float*d_xprev_init,float*d_v_first,
    int B,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)B*T, tot=BT*C;
    cublasHandle_t H; if(cublasCreate(&H)!=CUBLAS_STATUS_SUCCESS) return -1;
    tm_dw D; if(tm_upload_weights(P,&D)) return -1;

    /* device inputs */
    float *dxin,*dxprev=NULL,*dvf,*ddy,*d_vf_in=NULL;
    TCK(cudaMalloc(&dxin,tot*sizeof(float))); TCK(cudaMemcpy(dxin,x,tot*sizeof(float),cudaMemcpyHostToDevice));
    if(x_prev_init){ TCK(cudaMalloc(&dxprev,(size_t)B*C*sizeof(float)));
        TCK(cudaMemcpy(dxprev,x_prev_init,(size_t)B*C*sizeof(float),cudaMemcpyHostToDevice)); }
    TCK(cudaMalloc(&dvf,tot*sizeof(float)));
    if(L>0) TCK(cudaMemcpy(dvf,v_first,tot*sizeof(float),cudaMemcpyHostToDevice));
    TCK(cudaMalloc(&ddy,tot*sizeof(float))); TCK(cudaMemcpy(ddy,dy,tot*sizeof(float),cudaMemcpyHostToDevice));
    if(L==0 && d_v_first){ TCK(cudaMalloc(&d_vf_in,tot*sizeof(float)));
        TCK(cudaMemcpy(d_vf_in,d_v_first,tot*sizeof(float),cudaMemcpyHostToDevice)); }

    /* recompute forward to fill caches */
    tm_cache_d c; memset(&c,0,sizeof(c));
    float*dout_ign; TCK(cudaMalloc(&dout_ign,tot*sizeof(float)));
    if(tm_fwd_device(P,H,&D,dxin,dxprev,dvf,&c,dout_ign,B,T)) return -1;

    /* device grad mirror (all zeroed) */
    tm_dw Gd; memset(&Gd,0,sizeof(Gd));
    #define GA(f,n) do{ TCK(cudaMalloc(&Gd.f,(size_t)(n)*sizeof(float))); TCK(cudaMemset(Gd.f,0,(size_t)(n)*sizeof(float))); }while(0)
    GA(x_r,C);GA(x_k,C);GA(x_v,C);GA(x_w,C);GA(x_a,C);GA(x_g,C);
    GA(receptance_w,(size_t)C*C);GA(key_w,(size_t)C*C);GA(value_w,(size_t)C*C);GA(output_w,(size_t)C*C);
    GA(w0,C);GA(w1,(size_t)C*Dw);GA(w2,(size_t)Dw*C);
    GA(a0,C);GA(a1,(size_t)C*Da);GA(a2,(size_t)Da*C);
    GA(g1,(size_t)C*Dg);GA(g2,(size_t)Dg*C);
    if(L>0){GA(v0,C);GA(v1,(size_t)C*Dv);GA(v2,(size_t)Dv*C);}
    GA(k_k,C);GA(k_a,C);GA(r_k,C);GA(ln_x_w,C);GA(ln_x_b,C);
    #undef GA

    /* scratch */
    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, *d_vf_out=NULL;
    float *d_la,*d_tw,*d_lw,*d_sg,*d_lg,*d_lv=NULL;
    #define S(p,n) TCK(cudaMalloc(&p,(size_t)(n)*sizeof(float)))
    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_out,tot); TCK(cudaMemset(d_vf_out,0,tot*sizeof(float))); }
    #undef S

    /* ===== PHASE 2: post-processing backward ===== */
    /* output proj: dtg = dy @ Wo^T ; G_output_w += tg^T @ dy */
    TBK(rm_gemm(H,0,1,BT,C,C,1.0f,ddy,D.output_w,0.0f,dtg));
    TBK(rm_gemm(H,1,0,C,C,BT,1.0f,c.tg,ddy,0.0f,Gd.output_w));
    k_split<<<nblk(tot),TPB>>>(dtg,c.gvec,c.yb,dyb,dg_post,tot);
    k_bonus_bwd<<<nblk(BT*nH),TPB>>>(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<<<nblk(BT*nH),TPB>>>(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 backward ===== */
    { size_t stsz=(size_t)B*nH*N*N, sasz=(size_t)T*B*nH*N, scsz=(size_t)T*stsz;
      float*SC,*SA; TCK(cudaMalloc(&SC,scsz*sizeof(float))); TCK(cudaMalloc(&SA,sasz*sizeof(float)));
      size_t shf=((size_t)N*N+7*N)*sizeof(float);
      wkv7_fwd_kernel<<<B*nH,N,shf>>>(c.R,c.W,c.K,c.V,c.A,c.Bb,NULL,NULL,NULL,SC,SA,B,T,nH,N);
      size_t shb=(2*(size_t)N*N+9*N)*sizeof(float);
      wkv7_bwd_kernel<<<B*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,B,T,nH,N);
      cudaFree(SC); cudaFree(SA); }

    /* ===== PHASE 4: preprocessing backward ===== */
    k_p4_seed<<<nblk(tot),TPB>>>(gR,dr_post,gK,dkmod_post,gV,dvf_post,dr_t,dkmod_t,dvf_t,tot);
    k_comp_keymod<<<nblk(tot),TPB>>>(gA,gB,c.avec,c.kk,c.kvec,D.k_a,dkmod_t,dkk,dav,dkvec,Gd.k_a,C,tot);
    k_l2norm_bwd<<<nblk(BT*nH),TPB>>>(c.tk,dkk,dtk,nH,N,BT);
    k_tk_bwd<<<nblk(tot),TPB>>>(dtk,D.k_k,c.kvec,dkvec,Gd.k_k,C,tot);
    /* a gate */
    k_avec_bwd<<<nblk(tot),TPB>>>(dav,c.avec,d_oa,Gd.a0,C,tot);
    TBK(rm_gemm(H,0,1,BT,Da,C,1.0f,d_oa,D.a2,0.0f,d_la));
    TBK(rm_gemm(H,1,0,Da,C,BT,1.0f,c.la,d_oa,0.0f,Gd.a2));
    TBK(rm_gemm(H,0,1,BT,C,Da,1.0f,d_la,D.a1,0.0f,dsa));
    TBK(rm_gemm(H,1,0,C,Da,BT,1.0f,c.sha,d_la,0.0f,Gd.a1));
    /* w decay */
    k_wdecay_bwd<<<nblk(tot),TPB>>>(gW,c.W,c.sigw,d_ow,Gd.w0,C,tot);
    TBK(rm_gemm(H,0,1,BT,Dw,C,1.0f,d_ow,D.w2,0.0f,d_tw));
    TBK(rm_gemm(H,1,0,Dw,C,BT,1.0f,c.tw,d_ow,0.0f,Gd.w2));
    k_tanhderiv<<<nblk(BT*Dw),TPB>>>(d_tw,c.tw,d_lw,BT*Dw);
    TBK(rm_gemm(H,0,1,BT,C,Dw,1.0f,d_lw,D.w1,0.0f,dsw));
    TBK(rm_gemm(H,1,0,C,Dw,BT,1.0f,c.shw,d_lw,0.0f,Gd.w1));
    /* g gate */
    TBK(rm_gemm(H,0,1,BT,Dg,C,1.0f,dg_post,D.g2,0.0f,d_sg));
    TBK(rm_gemm(H,1,0,Dg,C,BT,1.0f,c.sg,dg_post,0.0f,Gd.g2));
    k_sigderiv<<<nblk(BT*Dg),TPB>>>(d_sg,c.sg,d_lg,BT*Dg);
    TBK(rm_gemm(H,0,1,BT,C,Dg,1.0f,d_lg,D.g1,0.0f,dsg));
    TBK(rm_gemm(H,1,0,C,Dg,BT,1.0f,c.shg,d_lg,0.0f,Gd.g1));
    /* value residual */
    if(L==0){ k_vvec0<<<nblk(tot),TPB>>>(dvf_t,d_vf_in,dvvec,tot); }
    else {
        k_vvec1<<<nblk(tot),TPB>>>(dvf_t,c.blend,dvf,c.vvec,D.v0,dvvec,d_vf_out,d_ov,Gd.v0,C,tot);
        TBK(rm_gemm(H,0,1,BT,Dv,C,1.0f,d_ov,D.v2,0.0f,d_lv));
        TBK(rm_gemm(H,1,0,Dv,C,BT,1.0f,c.lv,d_ov,0.0f,Gd.v2));
        TBK(rm_gemm(H,0,1,BT,C,Dv,1.0f,d_lv,D.v1,0.0f,dsv));   /* dsh_v (beta0) */
        TBK(rm_gemm(H,1,0,C,Dv,BT,1.0f,c.shv,d_lv,0.0f,Gd.v1));
    }
    /* projections r,k,v (dsh_v accumulates onto v1 contribution when L>0) */
    TBK(rm_gemm(H,0,1,BT,C,C,1.0f,dr_t,D.receptance_w,0.0f,dsr));
    TBK(rm_gemm(H,1,0,C,C,BT,1.0f,c.shr,dr_t,0.0f,Gd.receptance_w));
    TBK(rm_gemm(H,0,1,BT,C,C,1.0f,dkvec,D.key_w,0.0f,dsk));
    TBK(rm_gemm(H,1,0,C,C,BT,1.0f,c.shk,dkvec,0.0f,Gd.key_w));
    TBK(rm_gemm(H,0,1,BT,C,C,1.0f,dvvec,D.value_w,(L>0)?1.0f:0.0f,dsv));
    TBK(rm_gemm(H,1,0,C,C,BT,1.0f,c.shv,dvvec,0.0f,Gd.value_w));

    /* token-shift backward -> dx, d_xprev_init, G_x_* */
    float*ddxprev=NULL; if(d_xprev_init){ TCK(cudaMalloc(&ddxprev,(size_t)B*C*sizeof(float)));
        TCK(cudaMemset(ddxprev,0,(size_t)B*C*sizeof(float))); }
    float*ddx; TCK(cudaMalloc(&ddx,tot*sizeof(float)));
    k_tokenshift_bwd<<<nblk(tot),TPB>>>(dxin,dxprev,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,ddxprev,B,T,C);
    TCK(cudaDeviceSynchronize());

    /* ===== download ===== */
    #define DN(dev,host,n) do{ if(host && dev) TCK(cudaMemcpy(host,dev,(size_t)(n)*sizeof(float),cudaMemcpyDeviceToHost)); }while(0)
    TCK(cudaMemcpy(dx,ddx,tot*sizeof(float),cudaMemcpyDeviceToHost));
    if(d_xprev_init) DN(ddxprev,d_xprev_init,(size_t)B*C);
    if(L>0 && d_v_first) DN(d_vf_out,d_v_first,tot);
    DN(Gd.x_r,G->x_r,C);DN(Gd.x_k,G->x_k,C);DN(Gd.x_v,G->x_v,C);
    DN(Gd.x_w,G->x_w,C);DN(Gd.x_a,G->x_a,C);DN(Gd.x_g,G->x_g,C);
    DN(Gd.receptance_w,G->receptance_w,(size_t)C*C);DN(Gd.key_w,G->key_w,(size_t)C*C);
    DN(Gd.value_w,G->value_w,(size_t)C*C);DN(Gd.output_w,G->output_w,(size_t)C*C);
    DN(Gd.w0,G->w0,C);DN(Gd.w1,G->w1,(size_t)C*Dw);DN(Gd.w2,G->w2,(size_t)Dw*C);
    DN(Gd.a0,G->a0,C);DN(Gd.a1,G->a1,(size_t)C*Da);DN(Gd.a2,G->a2,(size_t)Da*C);
    DN(Gd.g1,G->g1,(size_t)C*Dg);DN(Gd.g2,G->g2,(size_t)Dg*C);
    if(L>0){DN(Gd.v0,G->v0,C);DN(Gd.v1,G->v1,(size_t)C*Dv);DN(Gd.v2,G->v2,(size_t)Dv*C);}
    DN(Gd.k_k,G->k_k,C);DN(Gd.k_a,G->k_a,C);DN(Gd.r_k,G->r_k,C);
    DN(Gd.ln_x_w,G->ln_x_w,C);DN(Gd.ln_x_b,G->ln_x_b,C);
    #undef DN

    /* free */
    tmc_free(&c); tm_free_weights(&D); tm_free_weights(&Gd);
    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_out,ddx,ddxprev,dout_ign,dxin,dxprev,dvf,ddy,d_vf_in};
    for(unsigned i=0;i<sizeof(sc)/sizeof(sc[0]);i++) if(sc[i]) cudaFree(sc[i]);
    cublasDestroy(H);
    return 0;
}