123 lines
6.2 KiB
Plaintext
123 lines
6.2 KiB
Plaintext
/* SPDX-License-Identifier: AGPL-3.0-or-later */
|
|
/* Copyright (C) 2026 David (Corey) Herrera */
|
|
/*
|
|
* test_moba_attn_cuda.cu — parity harness for the GPU MoBA sparse-attention
|
|
* core (cuda_kernels/moba_attn.cu) vs the CPU oracle rx_attn_score_av_forward /
|
|
* rx_attn_score_av_backward (rwkvx/rx_pretrain.c — the exact scores->softmax->AV
|
|
* and dQ/dK/dV routines rx_attn_forward_seq / rx_attn_backward_full use), on
|
|
* GPU 0. Fixed-seed generated q,k,v fixtures at the real config incl. N=64.
|
|
*
|
|
* Metric: max-abs-error relative to tensor scale (inf-norm) — NOT the |b|+eps
|
|
* metric, which is invalid for the signed near-zero reductions here. Bar 1e-3.
|
|
* Covers: dense fallback (moba_top_k=0), a genuine top-k sparse case, multi-head
|
|
* + batch, and clean-chunk (T%chunk_size==0) cases; plus one partial-chunk case.
|
|
* Checks attention output y AND dQ,dK,dV (and the softmax weights p).
|
|
*/
|
|
#include <cuda_runtime.h>
|
|
#include <cublas_v2.h>
|
|
#include <stdio.h>
|
|
#include <stdlib.h>
|
|
#include <stdint.h>
|
|
#include <string.h>
|
|
#include <math.h>
|
|
#include <time.h>
|
|
|
|
#include "cuda_kernels/moba_attn.cu"
|
|
|
|
extern "C" int rx_attn_score_av_forward(const rx_attn_params*,const float*,const float*,
|
|
const float*,int,int,float*,float*);
|
|
extern "C" int rx_attn_score_av_backward(const rx_attn_params*,const float*,const float*,
|
|
const float*,const float*,int,int,float*,float*,float*);
|
|
|
|
/* splitmix64 RNG (same as the timemix/wkv7 harnesses) */
|
|
static uint64_t sm_state;
|
|
static uint64_t sm_next(void){ uint64_t z=(sm_state+=0x9E3779B97F4A7C15ULL);
|
|
z=(z^(z>>30))*0xBF58476D1CE4E5B9ULL; z=(z^(z>>27))*0x94D049BB133111EBULL; return z^(z>>31); }
|
|
static float sm_uf(void){ return (float)((sm_next()>>11)*(1.0/9007199254740992.0)); }
|
|
static float sm_sym(float s){ return (2.0f*sm_uf()-1.0f)*s; }
|
|
static float* rnd(size_t n,float s){ float*p=(float*)malloc(n*4); for(size_t i=0;i<n;i++)p[i]=sm_sym(s); return p; }
|
|
|
|
static int g_fail=0;
|
|
static double chk(const char*tname,const float*gpu,const float*cpu,size_t n,double bar){
|
|
double maxabs=0.0,scale=0.0;
|
|
for(size_t i=0;i<n;i++){ double d=fabs((double)gpu[i]-(double)cpu[i]); if(d>maxabs)maxabs=d;
|
|
double c=fabs((double)cpu[i]); if(c>scale)scale=c; }
|
|
double rel=maxabs/(scale+1e-30); const char*st=(rel<bar)?"PASS":"FAIL"; if(rel>=bar)g_fail=1;
|
|
printf(" %-8s rel=%.3e maxabs=%.3e scale=%.3e [%s]\n",tname,rel,maxabs,scale,st);
|
|
return rel;
|
|
}
|
|
|
|
typedef struct { int B,T,H,N,cs,topk; } cfg;
|
|
|
|
static void run_case(const char*name, cfg cf, uint64_t seed, double bar){
|
|
int C=cf.H*cf.N; size_t tot=(size_t)cf.B*cf.T*C;
|
|
size_t nsc=(size_t)cf.B*cf.H*cf.T*cf.T;
|
|
sm_state=seed;
|
|
rx_attn_params P; memset(&P,0,sizeof(P));
|
|
P.c=C; P.n_head=cf.H; P.head_size=cf.N; P.chunk_size=cf.cs; P.moba_top_k=cf.topk;
|
|
|
|
float*q=rnd(tot,0.8f),*k=rnd(tot,0.8f),*v=rnd(tot,0.8f),*dy=rnd(tot,1.0f);
|
|
float*y_c=(float*)malloc(tot*4),*y_g=(float*)malloc(tot*4);
|
|
float*p_c=(float*)malloc(nsc*4),*p_g=(float*)malloc(nsc*4);
|
|
float*dq_c=(float*)malloc(tot*4),*dk_c=(float*)malloc(tot*4),*dv_c=(float*)malloc(tot*4);
|
|
float*dq_g=(float*)malloc(tot*4),*dk_g=(float*)malloc(tot*4),*dv_g=(float*)malloc(tot*4);
|
|
|
|
rx_attn_score_av_forward(&P,q,k,v,cf.B,cf.T,p_c,y_c);
|
|
rx_attn_score_av_backward(&P,q,k,v,dy,cf.B,cf.T,dq_c,dk_c,dv_c);
|
|
if(moba_attn_forward_cuda(&P,q,k,v,cf.B,cf.T,p_g,y_g)){printf(" %s GPU fwd ERR\n",name);g_fail=1;return;}
|
|
if(moba_attn_backward_cuda(&P,q,k,v,dy,cf.B,cf.T,dq_g,dk_g,dv_g)){printf(" %s GPU bwd ERR\n",name);g_fail=1;return;}
|
|
|
|
printf(" [%s] B=%d T=%d H=%d N=%d C=%d chunk=%d topk=%d bar=%.0e\n",
|
|
name,cf.B,cf.T,cf.H,cf.N,C,cf.cs,cf.topk,bar);
|
|
chk("y", y_g, y_c, tot, bar);
|
|
chk("p", p_g, p_c, nsc, bar);
|
|
chk("dQ", dq_g,dq_c,tot, bar);
|
|
chk("dK", dk_g,dk_c,tot, bar);
|
|
chk("dV", dv_g,dv_c,tot, bar);
|
|
|
|
free(q);free(k);free(v);free(dy);free(y_c);free(y_g);free(p_c);free(p_g);
|
|
free(dq_c);free(dk_c);free(dv_c);free(dq_g);free(dk_g);free(dv_g);
|
|
}
|
|
|
|
static void bench(cfg cf,int iters){
|
|
int C=cf.H*cf.N; size_t tot=(size_t)cf.B*cf.T*C, nsc=(size_t)cf.B*cf.H*cf.T*cf.T;
|
|
sm_state=99;
|
|
rx_attn_params P; memset(&P,0,sizeof(P));
|
|
P.c=C;P.n_head=cf.H;P.head_size=cf.N;P.chunk_size=cf.cs;P.moba_top_k=cf.topk;
|
|
float*q=rnd(tot,0.8f),*k=rnd(tot,0.8f),*v=rnd(tot,0.8f);
|
|
float*y=(float*)malloc(tot*4),*p=(float*)malloc(nsc*4);
|
|
moba_attn_forward_cuda(&P,q,k,v,cf.B,cf.T,p,y); /* warmup */
|
|
struct timespec t0,t1; clock_gettime(CLOCK_MONOTONIC,&t0);
|
|
for(int i=0;i<iters;i++) moba_attn_forward_cuda(&P,q,k,v,cf.B,cf.T,p,y);
|
|
clock_gettime(CLOCK_MONOTONIC,&t1);
|
|
double sec=(t1.tv_sec-t0.tv_sec)+(t1.tv_nsec-t0.tv_nsec)*1e-9;
|
|
double toks=(double)cf.B*cf.T*iters;
|
|
printf(" moba-attn fwd B=%d T=%d H=%d N=%d C=%d chunk=%d topk=%d: %d it in %.4fs -> %.0f tokens/sec (one K80 die; incl H2D+alloc)\n",
|
|
cf.B,cf.T,cf.H,cf.N,C,cf.cs,cf.topk,iters,sec,toks/sec);
|
|
free(q);free(k);free(v);free(y);free(p);
|
|
}
|
|
|
|
int main(void){
|
|
int dev=0; cudaSetDevice(dev); cudaDeviceProp pr; cudaGetDeviceProperties(&pr,dev);
|
|
printf("=== GPU MoBA sparse-attention core FORWARD+BACKWARD parity vs CPU oracle ===\n");
|
|
printf("device %d: %s cc %d.%d\n\n",dev,pr.name,pr.major,pr.minor);
|
|
|
|
/* dense fallback (moba_top_k=0 => plain causal), clean chunks */
|
|
run_case("dense s4", (cfg){.B=1,.T=8, .H=2,.N=4, .cs=2,.topk=0}, 11, 1e-3);
|
|
/* real head_dim N=64, dense fallback, clean chunks (T%cs==0) */
|
|
run_case("dense N64", (cfg){.B=1,.T=16,.H=2,.N=64,.cs=4,.topk=0}, 23, 1e-3);
|
|
/* genuine top-k sparse (drops chunks), N=64, clean chunks */
|
|
run_case("sparse N64",(cfg){.B=1,.T=16,.H=2,.N=64,.cs=4,.topk=2}, 37, 1e-3);
|
|
/* multi-head + batch, genuine sparse, clean chunks */
|
|
run_case("mh sparse", (cfg){.B=2,.T=12,.H=3,.N=64,.cs=4,.topk=1}, 41, 1e-3);
|
|
/* partial-chunk (T%cs!=0): the degenerate cs3/T7 config; tight bar expected
|
|
* since this is GPU-fp32 vs CPU-fp32 (NOT the fp32-vs-fp64 FD floor). */
|
|
run_case("partial", (cfg){.B=1,.T=7, .H=2,.N=4, .cs=3,.topk=0}, 53, 1e-3);
|
|
|
|
printf("\n=== throughput (moba-attn forward, one die) ===\n");
|
|
bench((cfg){.B=8,.T=512,.H=12,.N=64,.cs=64,.topk=0},20);
|
|
|
|
printf("\n%s\n", g_fail?"RESULT: FAIL":"RESULT: PASS");
|
|
return g_fail;
|
|
}
|