304 lines
14 KiB
Plaintext
304 lines
14 KiB
Plaintext
/* SPDX-License-Identifier: AGPL-3.0-or-later AND Apache-2.0 */
|
|
/* Copyright (C) 2026 David (Corey) Herrera */
|
|
/*
|
|
* Portions derived from wkv7_cuda_fp32.cu in RWKV-LM
|
|
* (https://github.com/BlinkDL/RWKV-LM), by Bo Peng (BlinkDL),
|
|
* licensed under the Apache License, Version 2.0. Modified by
|
|
* David (Corey) Herrera, 2026. See NOTICE and LICENSES/Apache-2.0.txt.
|
|
*/
|
|
/*
|
|
* wkv7.cu — RWKV-7 WKV7 forward + backward-through-time CUDA kernel (sm_37/K80).
|
|
*
|
|
* Per-head sequential scan (NOT a matmul): one CUDA block per (batch,head),
|
|
* N threads per block (thread i owns state row i), sequential over time.
|
|
* fp32 throughout. Reproduces the fp32 numerics and the exact grad contracts
|
|
* of the CPU oracle rx_wkv7_forward_seq / rx_wkv7_backward (rwkvx/rx_pretrain.c),
|
|
* which is a scalar-C transcription of BlinkDL Apache-2.0 wkv7_cuda_fp32.cu.
|
|
*
|
|
* Layouts (identical to the CPU oracle):
|
|
* input vectors r,w,k,v,a,b,dy,y : [B][T][H][N], idx = (((b*T+t)*H+h)*N+n)
|
|
* state S / S_cache[t] / dS : [B][H][N][N], idx = (((b*H+h)*N+i)*N+j)
|
|
* sa_cache[t] : [B][H][N], idx = ((b*H+h)*N+i)
|
|
* a = nkk = -kk ; b = kka = kk*iclr ; w = EFFECTIVE decay (dw is bare, §5).
|
|
*
|
|
* Recurrence (per head, per t):
|
|
* sa[i] = sum_j a[j]*S_prev[i][j]
|
|
* S[i][j] = S_prev[i][j]*w[j] + k[j]*v[i] + sa[i]*b[j]
|
|
* y[i] = sum_j S[i][j]*r[j]
|
|
*
|
|
* Reverse scan (carry dS = dL/dS_t, init 0):
|
|
* dr[j] = sum_i S_t[i][j]*dy[i]; dS[i][j] += dy[i]*r[j]
|
|
* dk[j] = sum_i dS[i][j]*v[i]; db[j] = sum_i dS[i][j]*sa[i]; dw[j] = sum_i dS[i][j]*S_prev[i][j]
|
|
* dv[i] = sum_j dS[i][j]*k[j]; dsa[i] = sum_j dS[i][j]*b[j]
|
|
* da[j] = sum_i dsa[i]*S_prev[i][j]
|
|
* dS[i][j] = dS[i][j]*w[j] + dsa[i]*a[j]
|
|
*
|
|
* Shared-mem budget (per block): fwd = (N*N + 7*N) floats. For the backward we
|
|
* keep ONLY two N*N tiles in shared — the carried dS and the S_prev tile read
|
|
* from the global forward trajectory — plus 9*N vectors, i.e. (2*N*N + 9*N)
|
|
* floats. S_t is never materialised: dr[j] = sum_i S_t[i][j]*dy[i] is expanded
|
|
* as w[j]*sum_i S_prev[i][j]*dy[i] + k[j]*sum_i v[i]*dy[i] + b[j]*sum_i sa[i]*dy[i]
|
|
* (division-free; all operands are S_prev + two scalars). At N=64 this is
|
|
* 2*4096 + 576 = 8768 floats = 34.25 KiB, well under the sm_37 48 KiB/block cap
|
|
* (the old 3*N*N layout was 51456 B and overflowed at N=64).
|
|
*/
|
|
#include <cuda_runtime.h>
|
|
#include <stdio.h>
|
|
|
|
/* ---------------- forward + trajectory cache ---------------- */
|
|
__global__ void wkv7_fwd_kernel(
|
|
const float *r, const float *w, const float *k, const float *v,
|
|
const float *a, const float *b, const float *s_init,
|
|
float *y, float *s_final, float *S_cache, float *sa_cache,
|
|
int B, int T, int H, int N)
|
|
{
|
|
int bh = blockIdx.x; /* 0..B*H-1 */
|
|
int bch = bh / H;
|
|
int h = bh % H;
|
|
int i = threadIdx.x; /* row */
|
|
if (i >= N) return;
|
|
|
|
extern __shared__ float sh[];
|
|
float *sS = sh; /* N*N */
|
|
float *sr = sS + N * N; /* N */
|
|
float *sw = sr + N;
|
|
float *sk = sw + N;
|
|
float *sv = sk + N;
|
|
float *sa = sv + N;
|
|
float *sb = sa + N;
|
|
|
|
for (int j = 0; j < N; j++)
|
|
sS[i * N + j] = s_init ? s_init[(((size_t)(bch * H + h) * N + i) * N) + j] : 0.0f;
|
|
__syncthreads();
|
|
|
|
size_t state_sz = (size_t)B * H * N * N;
|
|
for (int t = 0; t < T; t++) {
|
|
int idx = (((bch * T + t) * H + h) * N + i);
|
|
sr[i] = r[idx]; sw[i] = w[idx]; sk[i] = k[idx];
|
|
sv[i] = v[idx]; sa[i] = a[idx]; sb[i] = b[idx];
|
|
__syncthreads();
|
|
|
|
/* cache S_prev, compute sa[i] */
|
|
float saij = 0.0f;
|
|
float *Sc = S_cache + (size_t)t * state_sz;
|
|
size_t sbase = ((size_t)(bch * H + h) * N + i) * N;
|
|
for (int j = 0; j < N; j++) {
|
|
float sval = sS[i * N + j];
|
|
Sc[sbase + j] = sval;
|
|
saij += sa[j] * sval;
|
|
}
|
|
sa_cache[(size_t)t * B * H * N + (bch * H + h) * N + i] = saij;
|
|
|
|
/* update row i, read out y[i] */
|
|
float vi = sv[i];
|
|
float yy = 0.0f;
|
|
for (int j = 0; j < N; j++) {
|
|
float sval = sS[i * N + j] * sw[j] + sk[j] * vi + saij * sb[j];
|
|
sS[i * N + j] = sval;
|
|
yy += sval * sr[j];
|
|
}
|
|
if (y) y[idx] = yy;
|
|
__syncthreads();
|
|
}
|
|
if (s_final) {
|
|
size_t sbase = ((size_t)(bch * H + h) * N + i) * N;
|
|
for (int j = 0; j < N; j++) s_final[sbase + j] = sS[i * N + j];
|
|
}
|
|
}
|
|
|
|
/* ---------------- backward-through-time (reverse scan) ----------------
|
|
* Shared tiles: sdS (carried dS) and sSp (S_prev read from the global forward
|
|
* trajectory). S_t is NOT held in shared; dr is computed analytically from
|
|
* S_prev + two scalars (see file header). Budget: 2*N*N + 9*N floats. */
|
|
__global__ void wkv7_bwd_kernel(
|
|
const float *r, const float *w, const float *k, const float *v,
|
|
const float *a, const float *b, const float *dy,
|
|
const float *S_cache, const float *sa_cache,
|
|
float *dr, float *dw, float *dk, float *dv, float *da, float *db,
|
|
float *dS_out,
|
|
int B, int T, int H, int N)
|
|
{
|
|
int bh = blockIdx.x;
|
|
int bch = bh / H;
|
|
int h = bh % H;
|
|
int tid = threadIdx.x;
|
|
if (tid >= N) return;
|
|
|
|
extern __shared__ float sh[];
|
|
float *sdS = sh; /* N*N carried dS */
|
|
float *sSp = sdS + N * N; /* N*N S_prev (from global trajectory) */
|
|
float *sr = sSp + N * N; /* N each */
|
|
float *sw = sr + N;
|
|
float *sk = sw + N;
|
|
float *sv = sk + N;
|
|
float *saV = sv + N;
|
|
float *sb = saV + N;
|
|
float *sdy = sb + N;
|
|
float *ssa = sdy + N;
|
|
float *sdsa = ssa + N;
|
|
|
|
for (int j = 0; j < N; j++) sdS[tid * N + j] = 0.0f;
|
|
__syncthreads();
|
|
|
|
size_t state_sz = (size_t)B * H * N * N;
|
|
for (int t = T - 1; t >= 0; t--) {
|
|
int idx = (((bch * T + t) * H + h) * N + tid);
|
|
sr[tid] = r[idx]; sw[tid] = w[idx]; sk[tid] = k[idx]; sv[tid] = v[idx];
|
|
saV[tid] = a[idx]; sb[tid] = b[idx]; sdy[tid] = dy[idx];
|
|
ssa[tid] = sa_cache[(size_t)t * B * H * N + (bch * H + h) * N + tid];
|
|
const float *Sc = S_cache + (size_t)t * state_sz;
|
|
size_t sbase = ((size_t)(bch * H + h) * N + tid) * N;
|
|
for (int j = 0; j < N; j++) sSp[tid * N + j] = Sc[sbase + j];
|
|
__syncthreads();
|
|
|
|
/* (1) read-out grad dr[j] WITHOUT materialising S_t:
|
|
* dr[j] = sum_i S_t[i][j]*dy[i]
|
|
* = w[j]*sum_i S_prev[i][j]*dy[i]
|
|
* + k[j]*sum_i v[i]*dy[i]
|
|
* + b[j]*sum_i sa[i]*dy[i]
|
|
* (column tid of S_prev + two thread-invariant scalars; division-free).
|
|
* Then fold dy*r into the dS carry. */
|
|
float colSpDY = 0.0f, vDY = 0.0f, saDY = 0.0f;
|
|
for (int i = 0; i < N; i++) {
|
|
float dyi = sdy[i];
|
|
colSpDY += sSp[i * N + tid] * dyi; /* column tid of S_prev */
|
|
vDY += sv[i] * dyi;
|
|
saDY += ssa[i] * dyi;
|
|
}
|
|
dr[idx] = sw[tid] * colSpDY + sk[tid] * vDY + sb[tid] * saDY;
|
|
float dyt = sdy[tid];
|
|
for (int j = 0; j < N; j++) sdS[tid * N + j] += dyt * sr[j];
|
|
__syncthreads();
|
|
|
|
/* (2) column grads (dk,db,dw) sum over i ; row grads (dv,dsa) sum over j */
|
|
float dkj = 0.0f, dbj = 0.0f, dwj = 0.0f;
|
|
for (int i = 0; i < N; i++) {
|
|
float dSij = sdS[i * N + tid];
|
|
dkj += dSij * sv[i];
|
|
dbj += dSij * ssa[i];
|
|
dwj += dSij * sSp[i * N + tid];
|
|
}
|
|
dk[idx] = dkj; db[idx] = dbj; dw[idx] = dwj;
|
|
float dvi = 0.0f, dsai = 0.0f;
|
|
for (int j = 0; j < N; j++) {
|
|
float dSij = sdS[tid * N + j];
|
|
dvi += dSij * sk[j];
|
|
dsai += dSij * sb[j];
|
|
}
|
|
dv[idx] = dvi; sdsa[tid] = dsai;
|
|
__syncthreads();
|
|
|
|
/* (3) da[j] = sum_i dsa[i]*S_prev[i][j] ; dS carry -> t-1 */
|
|
float daj = 0.0f;
|
|
for (int i = 0; i < N; i++) daj += sdsa[i] * sSp[i * N + tid];
|
|
da[idx] = daj;
|
|
float dsai2 = sdsa[tid];
|
|
for (int j = 0; j < N; j++)
|
|
sdS[tid * N + j] = sdS[tid * N + j] * sw[j] + dsai2 * saV[j];
|
|
__syncthreads();
|
|
}
|
|
if (dS_out) {
|
|
size_t sbase = ((size_t)(bch * H + h) * N + tid) * N;
|
|
for (int j = 0; j < N; j++) dS_out[sbase + j] = sdS[tid * N + j];
|
|
}
|
|
}
|
|
|
|
/* ---------------- host launchers ---------------- */
|
|
#define CK(call) do { cudaError_t e_=(call); if(e_!=cudaSuccess){ \
|
|
fprintf(stderr,"CUDA error %s:%d: %s\n",__FILE__,__LINE__,cudaGetErrorString(e_)); \
|
|
return -1; } } while(0)
|
|
|
|
/* forward: fills y (may be NULL), and S_cache/sa_cache for the backward. */
|
|
extern "C" int wkv7_forward_cuda(
|
|
const float *r, const float *w, const float *k, const float *v,
|
|
const float *a, const float *b, const float *s_init,
|
|
float *y, float *s_final, float *S_cache_out, float *sa_cache_out,
|
|
int B, int T, int H, int N)
|
|
{
|
|
size_t nvec = (size_t)B * T * H * N;
|
|
size_t stsz = (size_t)B * H * N * N;
|
|
size_t sasz = (size_t)T * B * H * N;
|
|
size_t scsz = (size_t)T * stsz;
|
|
float *dr_,*dw_,*dk_,*dv_,*da_,*db_,*dsi=NULL,*dy_,*dsf=NULL,*dSc,*dsa;
|
|
CK(cudaMalloc(&dr_,nvec*sizeof(float))); CK(cudaMemcpy(dr_,r,nvec*sizeof(float),cudaMemcpyHostToDevice));
|
|
CK(cudaMalloc(&dw_,nvec*sizeof(float))); CK(cudaMemcpy(dw_,w,nvec*sizeof(float),cudaMemcpyHostToDevice));
|
|
CK(cudaMalloc(&dk_,nvec*sizeof(float))); CK(cudaMemcpy(dk_,k,nvec*sizeof(float),cudaMemcpyHostToDevice));
|
|
CK(cudaMalloc(&dv_,nvec*sizeof(float))); CK(cudaMemcpy(dv_,v,nvec*sizeof(float),cudaMemcpyHostToDevice));
|
|
CK(cudaMalloc(&da_,nvec*sizeof(float))); CK(cudaMemcpy(da_,a,nvec*sizeof(float),cudaMemcpyHostToDevice));
|
|
CK(cudaMalloc(&db_,nvec*sizeof(float))); CK(cudaMemcpy(db_,b,nvec*sizeof(float),cudaMemcpyHostToDevice));
|
|
if (s_init){ CK(cudaMalloc(&dsi,stsz*sizeof(float))); CK(cudaMemcpy(dsi,s_init,stsz*sizeof(float),cudaMemcpyHostToDevice)); }
|
|
CK(cudaMalloc(&dy_,nvec*sizeof(float)));
|
|
if (s_final) CK(cudaMalloc(&dsf,stsz*sizeof(float)));
|
|
CK(cudaMalloc(&dSc,scsz*sizeof(float)));
|
|
CK(cudaMalloc(&dsa,sasz*sizeof(float)));
|
|
|
|
size_t shmem = ((size_t)N*N + 7*N)*sizeof(float);
|
|
wkv7_fwd_kernel<<<B*H, N, shmem>>>(dr_,dw_,dk_,dv_,da_,db_,dsi,
|
|
y?dy_:NULL, dsf, dSc, dsa, B,T,H,N);
|
|
CK(cudaGetLastError()); CK(cudaDeviceSynchronize());
|
|
|
|
if (y) CK(cudaMemcpy(y,dy_,nvec*sizeof(float),cudaMemcpyDeviceToHost));
|
|
if (s_final) CK(cudaMemcpy(s_final,dsf,stsz*sizeof(float),cudaMemcpyDeviceToHost));
|
|
if (S_cache_out) CK(cudaMemcpy(S_cache_out,dSc,scsz*sizeof(float),cudaMemcpyDeviceToHost));
|
|
if (sa_cache_out) CK(cudaMemcpy(sa_cache_out,dsa,sasz*sizeof(float),cudaMemcpyDeviceToHost));
|
|
|
|
cudaFree(dr_);cudaFree(dw_);cudaFree(dk_);cudaFree(dv_);cudaFree(da_);cudaFree(db_);
|
|
if(dsi)cudaFree(dsi); cudaFree(dy_); if(dsf)cudaFree(dsf); cudaFree(dSc); cudaFree(dsa);
|
|
return 0;
|
|
}
|
|
|
|
extern "C" int wkv7_backward_cuda(
|
|
const float *r, const float *w, const float *k, const float *v,
|
|
const float *a, const float *b, const float *dy, const float *s_init,
|
|
float *dr, float *dw, float *dk, float *dv, float *da, float *db, float *dS_out,
|
|
int B, int T, int H, int N)
|
|
{
|
|
size_t nvec = (size_t)B * T * H * N;
|
|
size_t stsz = (size_t)B * H * N * N;
|
|
size_t sasz = (size_t)T * B * H * N;
|
|
size_t scsz = (size_t)T * stsz;
|
|
/* forward pass on device to fill caches */
|
|
float *dSc, *dsa;
|
|
CK(cudaMalloc(&dSc,scsz*sizeof(float)));
|
|
CK(cudaMalloc(&dsa,sasz*sizeof(float)));
|
|
/* device input copies */
|
|
float *R,*W,*K,*V,*A,*Bb,*DY,*Si=NULL;
|
|
CK(cudaMalloc(&R,nvec*sizeof(float))); CK(cudaMemcpy(R,r,nvec*sizeof(float),cudaMemcpyHostToDevice));
|
|
CK(cudaMalloc(&W,nvec*sizeof(float))); CK(cudaMemcpy(W,w,nvec*sizeof(float),cudaMemcpyHostToDevice));
|
|
CK(cudaMalloc(&K,nvec*sizeof(float))); CK(cudaMemcpy(K,k,nvec*sizeof(float),cudaMemcpyHostToDevice));
|
|
CK(cudaMalloc(&V,nvec*sizeof(float))); CK(cudaMemcpy(V,v,nvec*sizeof(float),cudaMemcpyHostToDevice));
|
|
CK(cudaMalloc(&A,nvec*sizeof(float))); CK(cudaMemcpy(A,a,nvec*sizeof(float),cudaMemcpyHostToDevice));
|
|
CK(cudaMalloc(&Bb,nvec*sizeof(float))); CK(cudaMemcpy(Bb,b,nvec*sizeof(float),cudaMemcpyHostToDevice));
|
|
CK(cudaMalloc(&DY,nvec*sizeof(float))); CK(cudaMemcpy(DY,dy,nvec*sizeof(float),cudaMemcpyHostToDevice));
|
|
if (s_init){ CK(cudaMalloc(&Si,stsz*sizeof(float))); CK(cudaMemcpy(Si,s_init,stsz*sizeof(float),cudaMemcpyHostToDevice)); }
|
|
|
|
size_t shf = ((size_t)N*N + 7*N)*sizeof(float);
|
|
wkv7_fwd_kernel<<<B*H, N, shf>>>(R,W,K,V,A,Bb,Si,NULL,NULL,dSc,dsa,B,T,H,N);
|
|
CK(cudaGetLastError()); CK(cudaDeviceSynchronize());
|
|
|
|
/* backward grads */
|
|
float *dR,*dW,*dK,*dV,*dA,*dB,*dSo;
|
|
CK(cudaMalloc(&dR,nvec*sizeof(float))); CK(cudaMalloc(&dW,nvec*sizeof(float)));
|
|
CK(cudaMalloc(&dK,nvec*sizeof(float))); CK(cudaMalloc(&dV,nvec*sizeof(float)));
|
|
CK(cudaMalloc(&dA,nvec*sizeof(float))); CK(cudaMalloc(&dB,nvec*sizeof(float)));
|
|
CK(cudaMalloc(&dSo,stsz*sizeof(float)));
|
|
|
|
size_t shb = (2*(size_t)N*N + 9*N)*sizeof(float);
|
|
wkv7_bwd_kernel<<<B*H, N, shb>>>(R,W,K,V,A,Bb,DY,dSc,dsa,
|
|
dR,dW,dK,dV,dA,dB,dSo,B,T,H,N);
|
|
CK(cudaGetLastError()); CK(cudaDeviceSynchronize());
|
|
|
|
if(dr)CK(cudaMemcpy(dr,dR,nvec*sizeof(float),cudaMemcpyDeviceToHost));
|
|
if(dw)CK(cudaMemcpy(dw,dW,nvec*sizeof(float),cudaMemcpyDeviceToHost));
|
|
if(dk)CK(cudaMemcpy(dk,dK,nvec*sizeof(float),cudaMemcpyDeviceToHost));
|
|
if(dv)CK(cudaMemcpy(dv,dV,nvec*sizeof(float),cudaMemcpyDeviceToHost));
|
|
if(da)CK(cudaMemcpy(da,dA,nvec*sizeof(float),cudaMemcpyDeviceToHost));
|
|
if(db)CK(cudaMemcpy(db,dB,nvec*sizeof(float),cudaMemcpyDeviceToHost));
|
|
if(dS_out)CK(cudaMemcpy(dS_out,dSo,stsz*sizeof(float),cudaMemcpyDeviceToHost));
|
|
|
|
cudaFree(dSc);cudaFree(dsa);cudaFree(R);cudaFree(W);cudaFree(K);cudaFree(V);
|
|
cudaFree(A);cudaFree(Bb);cudaFree(DY); if(Si)cudaFree(Si);
|
|
cudaFree(dR);cudaFree(dW);cudaFree(dK);cudaFree(dV);cudaFree(dA);cudaFree(dB);cudaFree(dSo);
|
|
return 0;
|
|
}
|