/* dragon_gpu.cu — Phase 5a: attention weights RESIDENT in VRAM, bf16 matvec on GPU.
 * Strategy (profile-driven): the box is DISK-bound (48GB/token working set > 31GB RAM).
 * VRAM = a permanent 3rd cache tier: park attention weights (per sub-block ~181MB bf16) for as
 * many of the 56 sub-blocks as fit (~7.5GB → ~40 blocks). Two compound wins:
 *   1) attention matvecs run on GPU (fast, weights never leave VRAM)
 *   2) ~7.5GB leaves the RAM working set → page cache holds more dense/experts → fewer disk reads
 * Experts stay on CPU (the advisor's trap-note is right: never stream 12GB/token over PCIe).
 * Build: nvcc -O3 -arch=sm_89 -Xcompiler -fopenmp -c dragon_gpu.cu -o dragon_gpu.o
 */
#include <cuda_runtime.h>
#include <cuda_bf16.h>
#include <stdio.h>
#include <stdint.h>

#define CK(x) do{cudaError_t e=(x); if(e!=cudaSuccess){fprintf(stderr,"CUDA %s:%d %s\n",__FILE__,__LINE__,cudaGetErrorString(e)); exit(9);} }while(0)

/* one warp per output row: o[i] = sum_k bf16(W[i,k]) * x[k] */
__global__ void mv_bf16_kernel(const __nv_bfloat16* __restrict__ W, const float* __restrict__ x,
                               float* __restrict__ o, int M, int K){
    int row = blockIdx.x*(blockDim.x/32) + (threadIdx.x/32);
    int lane = threadIdx.x & 31;
    if(row>=M) return;
    const __nv_bfloat16* w = W + (long)row*K;
    float a = 0.f;
    for(int k=lane; k<K; k+=32) a += __bfloat162float(w[k]) * x[k];
    for(int off=16; off>0; off>>=1) a += __shfl_down_sync(0xffffffff,a,off);
    if(lane==0) o[row]=a;
}

extern "C" {
static float *d_x, *d_o;          /* device activation buffers */
static void* d_attn[56][5];       /* resident attention weights per sub-block: qa,qb,kva,kvb,op */
static int   resident[56];        /* 1 if this sub-block's attention lives in VRAM */
static const int MDIM[5]={1536,12288,576,16384,6144};
static const int KDIM[5]={6144,1536,6144,512,8192};

int gpu_init(int max_blocks){
    if(cudaSetDevice(0)!=cudaSuccess) return 0;
    size_t freeb,totalb; cudaMemGetInfo(&freeb,&totalb);
    fprintf(stderr,"[gpu] free %zuMB / %zuMB\n",freeb>>20,totalb>>20);
    CK(cudaMalloc(&d_x, 16384*sizeof(float)));
    CK(cudaMalloc(&d_o, 16384*sizeof(float)));
    return 1;
}
/* upload one sub-block's 5 attention weights (host bf16 pointers straight from mmap) */
int gpu_upload_attn(int sb, const void* qa,const void* qb,const void* kva,const void* kvb,const void* op){
    const void* src[5]={qa,qb,kva,kvb,op};
    size_t need=0; for(int i=0;i<5;i++) need += (size_t)MDIM[i]*KDIM[i]*2;
    size_t freeb,totalb; cudaMemGetInfo(&freeb,&totalb);
    if(freeb < need + (200u<<20)) return 0;   /* keep 200MB headroom */
    for(int i=0;i<5;i++){
        size_t nb=(size_t)MDIM[i]*KDIM[i]*2;
        CK(cudaMalloc(&d_attn[sb][i],nb));
        CK(cudaMemcpy(d_attn[sb][i],src[i],nb,cudaMemcpyHostToDevice));
    }
    resident[sb]=1; return 1;
}
int gpu_has(int sb){ return resident[sb]; }
/* o[M] = W_idx @ x[K] on GPU for resident sub-block sb, weight idx w (0..4) */
void gpu_mv(int sb, int wi, const float* x, float* o){
    int M=MDIM[wi], K=KDIM[wi];
    CK(cudaMemcpy(d_x,x,K*sizeof(float),cudaMemcpyHostToDevice));
    int warps_per_block=8, threads=warps_per_block*32;
    int blocks=(M+warps_per_block-1)/warps_per_block;
    mv_bf16_kernel<<<blocks,threads>>>((const __nv_bfloat16*)d_attn[sb][wi],d_x,d_o,M,K);
    CK(cudaMemcpy(o,d_o,M*sizeof(float),cudaMemcpyDeviceToHost));
}
} /* extern C */