276 lines
14 KiB
Plaintext
276 lines
14 KiB
Plaintext
/* 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;
|
||
}
|