/* SPDX-License-Identifier: AGPL-3.0-or-later */
/* Copyright (C) 2026 David (Corey) Herrera */
/*
 * moba_attn.cu — GPU MoBA (Mixture-of-Block-Attention) sparse causal
 * softmax-attention core for RWKV-x pretraining, forward + backward, sm_37/K80.
 *
 * This is the ISOLATED attention core (the "last hard kernel"): given the
 * already-projected per-token q,k,v [B*T*C] (C=n_head*head_size), it computes
 *   scores = scale * Q·Kᵀ  ->  MoBA-masked causal softmax  ->  Y = A·V
 * and the full backward  d(Y) -> dQ,dK,dV.  Parity target is the CPU oracle
 * attn_score_av / attn_score_av_bwd, exported from rwkvx/rx_pretrain.c as
 * rx_attn_score_av_forward / rx_attn_score_av_backward (the exact routines
 * rx_attn_forward_seq / rx_attn_backward_full use).
 *
 * Layout / routing (mirrors the timemix.cu conventions):
 *   - q,k,v arrive [B,T,C] (host).  A pack kernel transposes them to per-head
 *     contiguous qh,kh,vh [B,H,T,N] so each (b,h) slice is a packed row-major
 *     [T,N] block; results yh/dqh/dkh/dvh unpack back to [B,T,C].
 *   - Q·Kᵀ, A·V and all four backward gemms (dP, dV, dQ, dK) go through the
 *     rm_gemm() row-major cuBLAS helper (swap A/B; N,M,K), one call per (b,h).
 *     The score gemm folds the 1/sqrt(N) scale into alpha.
 *   - Block-selection, softmax and the softmax backward are per-query-row
 *     reduction kernels (one thread per (b,h,tq)).  The softmax normalizer and
 *     the backward score-gradient reduction (wsum) are DOUBLE-accumulated —
 *     cancellation-prone, same discipline as the LN backward.
 *
 * MoBA selection on GPU: the per-row kernel recomputes the gate g_c =
 *   q · mean_{s∈chunk c}(k)  in DOUBLE, in the exact same order as the oracle,
 *   from the same fp32 q,k — so the top-k argmax selection is bit-identical to
 *   the CPU side (no near-tie divergence).  The selection is non-differentiable
 *   and treated as FIXED: unselected/non-causal keys get softmax weight exactly
 *   0.  Those zeros in P then carry the mask through EVERY backward gemm with no
 *   extra masking — dS = scale·P·(dP−wsum) and dV = Pᵀ·dY both vanish wherever
 *   P=0 — so only the forward softmax needs the mask.  moba_top_k<=0 (or
 *   >=#candidate chunks for a given query) is the dense/all-select fallback =
 *   plain causal attention.
 *
 * Numerics: GPU fp32 (cuBLAS Sgemm) vs the oracle's fp64-accumulated scores;
 * expect ~1e-6..1e-5 rel, well under the 1e-3 inf-norm bar.
 *
 * Compile (K80 host, standalone — NOT the fc_cuda_ptx.h embed pipeline):
 *   nvcc -arch=sm_37 -ccbin gcc-11 -isystem cuda_kernels/nvcc_shim -I. ...
 */
#include <cuda_runtime.h>
#include <cublas_v2.h>
#include <stdio.h>
#include <stdlib.h>
#include <math.h>

#include "rwkvx/rx_pretrain.h"   /* rx_attn_params */
#include "cuda_kernels/fm_arena.cuh" /* device scratch bump-arena (opt #2) */

#define MOBA_TPB 256
static inline int moba_nblk(size_t n){ return (int)((n+MOBA_TPB-1)/MOBA_TPB); }

/* max #candidate chunks per query supported by the sparse-path per-row kernel
 * (ceil(T/chunk_size) <= this); host asserts it. */
#define MOBA_MAXCAND 512

#define MCK(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 MCB(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)

/* ---- row-major GEMM via cuBLAS: C[M,N] = alpha*op(A)@op(B) + beta*C ----
 * identical to timemix.cu / fc_cuda.c fc_sgemm_ctx: compute Cᵀ = op(B)ᵀ@op(A)ᵀ
 * by swapping A/B and passing (N,M,K). Packed row-major leading dims. */
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);
}

/* ============================ pack / unpack ============================ */
/* [B,T,C=H*N] <-> [B,H,T,N] (per-head contiguous). */
__global__ void moba_k_pack(const float*src,float*dst,int B,int T,int H,int N){
    size_t i=(size_t)blockIdx.x*blockDim.x+threadIdx.x; size_t tot=(size_t)B*T*H*N;
    if(i>=tot) return;
    int C=H*N; int c=(int)(i%C); size_t bt=i/C; int t=(int)(bt%T); int b=(int)(bt/T);
    int h=c/N, n=c%N;
    dst[(((size_t)(b*H+h))*T+t)*N+n]=src[i];
}
__global__ void moba_k_unpack(const float*src,float*dst,int B,int T,int H,int N){
    size_t i=(size_t)blockIdx.x*blockDim.x+threadIdx.x; size_t tot=(size_t)B*T*H*N;
    if(i>=tot) return;
    int C=H*N; int c=(int)(i%C); size_t bt=i/C; int t=(int)(bt%T); int b=(int)(bt/T);
    int h=c/N, n=c%N;
    dst[i]=src[(((size_t)(b*H+h))*T+t)*N+n];
}

/* ============================ forward softmax ============================ */
/* Dense/all-select fallback = plain causal softmax over s in [0,tq].
 * One thread per (b,h,tq); double normalizer; ascending-s order (= oracle). */
__global__ void moba_k_softmax_causal(const float*S,float*P,long BH,int T){
    size_t r=(size_t)blockIdx.x*blockDim.x+threadIdx.x; if(r>=(size_t)BH*T) return;
    int tq=(int)(r%T); const float*Sr=S+r*T; float*Pr=P+r*T;
    double maxs=-1e300;
    for(int s=0;s<=tq;s++){ double d=(double)Sr[s]; if(d>maxs)maxs=d; }
    double sum=0.0;
    for(int s=0;s<=tq;s++) sum+=exp((double)Sr[s]-maxs);
    for(int s=0;s<T;s++)
        Pr[s]= (s<=tq)? (float)(exp((double)Sr[s]-maxs)/sum) : 0.0f;
}

/* True MoBA-sparse softmax: recompute gate selection (double, oracle order)
 * then softmax over the included causal keys.  Per (b,h,tq). */
__global__ void moba_k_softmax_moba(const float*S,const float*qh,const float*kh,
    float*P,int B,int H,int T,int N,int cs,int topk){
    size_t r=(size_t)blockIdx.x*blockDim.x+threadIdx.x; if(r>=(size_t)B*H*T) return;
    int tq=(int)(r%T); size_t bh=r/T;
    const float*Sr=S+r*T; float*Pr=P+r*T;
    const float*qrow=qh + (bh*T+tq)*N;
    const float*kbase=kh + bh*T*N;               /* [T,N] for this head */
    int local_chunk=tq/cs, ncand=local_chunk;
    unsigned char csel[MOBA_MAXCAND];
    int keepall=(topk<=0 || topk>=ncand);
    if(keepall){ for(int c=0;c<ncand;c++) csel[c]=1; }
    else {
        double gate[MOBA_MAXCAND];
        for(int c=0;c<ncand;c++){ double g=0.0;
            for(int j=0;j<N;j++){ double mkj=0.0;
                for(int s=c*cs;s<(c+1)*cs;s++) mkj+=(double)kbase[(size_t)s*N+j];
                mkj/=(double)cs; g+=(double)qrow[j]*mkj; }
            gate[c]=g; csel[c]=0; }
        for(int it=0;it<topk;it++){ int arg=-1; double best=-1e300;
            for(int c=0;c<ncand;c++) if(!csel[c]&&gate[c]>best){best=gate[c];arg=c;}
            if(arg<0)break; csel[arg]=1; }
    }
    double maxs=-1e300;
    for(int s=0;s<=tq;s++){ int c=s/cs; int inc=(c==local_chunk)||(c<ncand&&csel[c]);
        if(inc){ double d=(double)Sr[s]; if(d>maxs)maxs=d; } }
    double sum=0.0;
    for(int s=0;s<=tq;s++){ int c=s/cs; int inc=(c==local_chunk)||(c<ncand&&csel[c]);
        if(inc) sum+=exp((double)Sr[s]-maxs); }
    for(int s=0;s<T;s++){ int c=s/cs;
        int inc=(s<=tq)&&((c==local_chunk)||(c<ncand&&csel[c]));
        Pr[s]= inc? (float)(exp((double)Sr[s]-maxs)/sum) : 0.0f; }
}

/* ============================ backward softmax ============================ */
/* dS[tq,s] = scale * P[tq,s] * (dP[tq,s] - wsum),  wsum = Σ_s P·dP (double).
 * P's zeros carry the MoBA/causal mask; nothing else to mask. */
__global__ void moba_k_softmax_bwd(const float*P,const float*dP,float*dS,
    long BH,int T,float scale){
    size_t r=(size_t)blockIdx.x*blockDim.x+threadIdx.x; if(r>=(size_t)BH*T) return;
    int tq=(int)(r%T); const float*Pr=P+r*T,*dPr=dP+r*T; float*dSr=dS+r*T;
    double wsum=0.0;
    for(int s=0;s<=tq;s++){ double pv=(double)Pr[s]; if(pv!=0.0) wsum+=pv*(double)dPr[s]; }
    for(int s=0;s<T;s++){ double pv=(double)Pr[s];
        dSr[s]= (s<=tq && pv!=0.0)? (float)((double)scale*pv*((double)dPr[s]-wsum)) : 0.0f; }
}

/* ============================ device driver ============================ */
typedef struct { float *qh,*kh,*vh,*S,*P; } moba_dev;

static int moba_alloc(moba_dev*d,size_t tot,size_t nsc){
    d->qh=(float*)fmA_alloc(tot*sizeof(float));
    d->kh=(float*)fmA_alloc(tot*sizeof(float));
    d->vh=(float*)fmA_alloc(tot*sizeof(float));
    d->S =(float*)fmA_alloc(nsc*sizeof(float));
    d->P =(float*)fmA_alloc(nsc*sizeof(float));
    return 0;
}
static void moba_free(moba_dev*d){
    fmA_free(d->qh);fmA_free(d->kh);fmA_free(d->vh);fmA_free(d->S);fmA_free(d->P);
}

/* Compute P (softmax weights) from packed qh,kh; leaves S,P on device. */
static int moba_scores_softmax(cublasHandle_t H,moba_dev*d,int B,int nH,int T,int N,
    int cs,int topk,float scale){
    long BH=(long)B*nH;
    for(long bh=0;bh<BH;bh++)
        MCB(rm_gemm(H,0,1,T,T,N,scale, d->qh+(size_t)bh*T*N, d->kh+(size_t)bh*T*N,
                    0.0f, d->S+(size_t)bh*T*T));
    if(topk<=0)
        moba_k_softmax_causal<<<moba_nblk((size_t)BH*T),MOBA_TPB>>>(d->S,d->P,BH,T);
    else
        moba_k_softmax_moba<<<moba_nblk((size_t)BH*T),MOBA_TPB>>>(d->S,d->qh,d->kh,d->P,
            B,nH,T,N,cs,topk);
    return 0;
}

extern "C" int moba_attn_forward_cuda(const rx_attn_params*P,const float*q,
    const float*k,const float*v,int B,int T,float*p_out,float*y_out)
{
    int C=P->c,nH=P->n_head,N=P->head_size,cs=P->chunk_size<1?1:P->chunk_size;
    int topk=P->moba_top_k; float scale=1.0f/sqrtf((float)N);
    size_t tot=(size_t)B*T*C, nsc=(size_t)B*nH*T*T;
    if((T+cs-1)/cs>MOBA_MAXCAND){ fprintf(stderr,"moba: ncand>%d\n",MOBA_MAXCAND); return -1; }

    moba_dev d; if(moba_alloc(&d,tot,nsc)) return -1;
    float *gq,*gk,*gv,*gy,*gYh;
    MCK(cudaMalloc(&gq,tot*sizeof(float))); MCK(cudaMalloc(&gk,tot*sizeof(float)));
    MCK(cudaMalloc(&gv,tot*sizeof(float))); MCK(cudaMalloc(&gy,tot*sizeof(float)));
    MCK(cudaMalloc(&gYh,tot*sizeof(float)));
    MCK(cudaMemcpy(gq,q,tot*sizeof(float),cudaMemcpyHostToDevice));
    MCK(cudaMemcpy(gk,k,tot*sizeof(float),cudaMemcpyHostToDevice));
    MCK(cudaMemcpy(gv,v,tot*sizeof(float),cudaMemcpyHostToDevice));
    cublasHandle_t H; if(cublasCreate(&H)!=CUBLAS_STATUS_SUCCESS){moba_free(&d);return -1;}

    int nb=moba_nblk(tot);
    moba_k_pack<<<nb,MOBA_TPB>>>(gq,d.qh,B,T,nH,N);
    moba_k_pack<<<nb,MOBA_TPB>>>(gk,d.kh,B,T,nH,N);
    moba_k_pack<<<nb,MOBA_TPB>>>(gv,d.vh,B,T,nH,N);
    if(moba_scores_softmax(H,&d,B,nH,T,N,cs,topk,scale)){cublasDestroy(H);moba_free(&d);return -1;}
    long BH=(long)B*nH;
    for(long bh=0;bh<BH;bh++)               /* Y = P · vh */
        MCB(rm_gemm(H,0,0,T,N,T,1.0f, d.P+(size_t)bh*T*T, d.vh+(size_t)bh*T*N,
                    0.0f, gYh+(size_t)bh*T*N));
    moba_k_unpack<<<nb,MOBA_TPB>>>(gYh,gy,B,T,nH,N);

    MCK(cudaMemcpy(y_out,gy,tot*sizeof(float),cudaMemcpyDeviceToHost));
    if(p_out) MCK(cudaMemcpy(p_out,d.P,nsc*sizeof(float),cudaMemcpyDeviceToHost));
    MCK(cudaDeviceSynchronize());
    cublasDestroy(H); moba_free(&d);
    cudaFree(gq);cudaFree(gk);cudaFree(gv);cudaFree(gy);cudaFree(gYh);
    return 0;
}

extern "C" int moba_attn_backward_cuda(const rx_attn_params*P,const float*q,
    const float*k,const float*v,const float*dy,int B,int T,
    float*dq_out,float*dk_out,float*dv_out)
{
    int C=P->c,nH=P->n_head,N=P->head_size,cs=P->chunk_size<1?1:P->chunk_size;
    int topk=P->moba_top_k; float scale=1.0f/sqrtf((float)N);
    size_t tot=(size_t)B*T*C, nsc=(size_t)B*nH*T*T;
    if((T+cs-1)/cs>MOBA_MAXCAND){ fprintf(stderr,"moba: ncand>%d\n",MOBA_MAXCAND); return -1; }

    moba_dev d; if(moba_alloc(&d,tot,nsc)) return -1;
    float *gq,*gk,*gv,*gdy,*dyh,*dP,*dS,*dqh,*dkh,*dvh,*gdq,*gdk,*gdv;
    #define M_(p) MCK(cudaMalloc(&p,tot*sizeof(float)))
    M_(gq);M_(gk);M_(gv);M_(gdy);M_(dyh);M_(dqh);M_(dkh);M_(dvh);M_(gdq);M_(gdk);M_(gdv);
    #undef M_
    MCK(cudaMalloc(&dP,nsc*sizeof(float))); MCK(cudaMalloc(&dS,nsc*sizeof(float)));
    MCK(cudaMemcpy(gq ,q ,tot*sizeof(float),cudaMemcpyHostToDevice));
    MCK(cudaMemcpy(gk ,k ,tot*sizeof(float),cudaMemcpyHostToDevice));
    MCK(cudaMemcpy(gv ,v ,tot*sizeof(float),cudaMemcpyHostToDevice));
    MCK(cudaMemcpy(gdy,dy,tot*sizeof(float),cudaMemcpyHostToDevice));
    cublasHandle_t H; if(cublasCreate(&H)!=CUBLAS_STATUS_SUCCESS){moba_free(&d);return -1;}

    int nb=moba_nblk(tot); long BH=(long)B*nH;
    moba_k_pack<<<nb,MOBA_TPB>>>(gq ,d.qh,B,T,nH,N);
    moba_k_pack<<<nb,MOBA_TPB>>>(gk ,d.kh,B,T,nH,N);
    moba_k_pack<<<nb,MOBA_TPB>>>(gv ,d.vh,B,T,nH,N);
    moba_k_pack<<<nb,MOBA_TPB>>>(gdy,dyh ,B,T,nH,N);
    /* recompute selection + softmax weights P (fixed for the backward) */
    if(moba_scores_softmax(H,&d,B,nH,T,N,cs,topk,scale)){cublasDestroy(H);moba_free(&d);return -1;}

    for(long bh=0;bh<BH;bh++){
        size_t os=(size_t)bh*T*T, ov=(size_t)bh*T*N;
        MCB(rm_gemm(H,0,1,T,T,N,1.0f, dyh+ov, d.vh+ov, 0.0f, dP+os));  /* dP = dY·Vᵀ */
        MCB(rm_gemm(H,1,0,T,N,T,1.0f, d.P+os, dyh+ov, 0.0f, dvh+ov));  /* dV = Pᵀ·dY */
    }
    moba_k_softmax_bwd<<<moba_nblk((size_t)BH*T),MOBA_TPB>>>(d.P,dP,dS,BH,T,scale);
    for(long bh=0;bh<BH;bh++){
        size_t os=(size_t)bh*T*T, ov=(size_t)bh*T*N;
        MCB(rm_gemm(H,0,0,T,N,T,1.0f, dS+os, d.kh+ov, 0.0f, dqh+ov));  /* dQ = dS·K  */
        MCB(rm_gemm(H,1,0,T,N,T,1.0f, dS+os, d.qh+ov, 0.0f, dkh+ov));  /* dK = dSᵀ·Q */
    }
    moba_k_unpack<<<nb,MOBA_TPB>>>(dqh,gdq,B,T,nH,N);
    moba_k_unpack<<<nb,MOBA_TPB>>>(dkh,gdk,B,T,nH,N);
    moba_k_unpack<<<nb,MOBA_TPB>>>(dvh,gdv,B,T,nH,N);

    MCK(cudaMemcpy(dq_out,gdq,tot*sizeof(float),cudaMemcpyDeviceToHost));
    MCK(cudaMemcpy(dk_out,gdk,tot*sizeof(float),cudaMemcpyDeviceToHost));
    MCK(cudaMemcpy(dv_out,gdv,tot*sizeof(float),cudaMemcpyDeviceToHost));
    MCK(cudaDeviceSynchronize());
    cublasDestroy(H); moba_free(&d);
    cudaFree(gq);cudaFree(gk);cudaFree(gv);cudaFree(gdy);cudaFree(dyh);
    cudaFree(dP);cudaFree(dS);cudaFree(dqh);cudaFree(dkh);cudaFree(dvh);
    cudaFree(gdq);cudaFree(gdk);cudaFree(gdv);
    return 0;
}