code wiki / _hdl_build / nx_intfp_lm_train_gate.nx
nx_intfp_lm_train_gate.nx
buildroot/runtime/_hdl_build/nx_intfp_lm_train_gate.nx
about
nx_intfp_lm_train_gate.nx -- COMPLETE integer transformer LM, trained on next-token prediction. Composes the
proven pieces: token EMBEDDING (lookup + scatter grad) -> full transformer LAYER (attn+FFN+2 norms+2 residuals,
gradchecked 84/84 elsewhere) -> output HEAD (y@Wout) -> CROSS-ENTROPY (softmax + the exact p-onehot gradient).
Everything Q20 integer, no float. Trains tokens [1,2,3,4,5,6] -> targets [2,3,4,5,6,7]; reports the (1-p_correct)
proxy loss + argmax ACCURACY. If accuracy -> T/T, a sovereign integer transformer LM has LEARNED. This is the
sovereign half of the PyTorch training h2h; only the run itself justifies a speed claim. license_tier: ORIGINAL
dependencies 1 imports · 0 importers
imports: nx_syscalls.nx
imported by: nobody (leaf or entry point)
call flow from main pre-order; caps 40 nodes / depth 6 declared; ↻ = already shown
structs
| none |
consts
| 12 | const S: i64 = 1048576 |
| 13 | const T: i64 = 6 |
| 14 | const DM: i64 = 4 |
| 15 | const HF: i64 = 8 |
| 16 | const VOC: i64 = 8 |
| 17 | const SCALE: i64 = 524288 |
| 18 | const EPS: i64 = 1048576 |
functions
| 9 | func w(s: *u8) -> i64 { var n: i64=0; while s[n]!=(0 as u8){n=n+1} sys_write(1,s,n); return 0 } |
| 10 | func wn(v: i64) -> i64 { if v==0 { sys_write(1,"0" as *u8,1); return 0 } var m: i64=v; if m<0{sys_write(1,"-" as *u8,1);m=0-m} let t: *u8=sys_mmap(24); var k: i64=0; while m>0{t[k]=(48+(m%10)) as u8;m=m/10;k=k+1} let o: *u8=sys_mmap(24); var q: i64=k-1; var i: i64=0; while q>=0{o[i]=t[q];i=i+1;q=q-1} sys_write(1,o,i); return 0 } |
| 20 | func isqrt(n: i64) -> i64 { if n<=0 { return 0 } var bit: i64=1; while bit*4<=n { bit=bit*4 } var res: i64=0; var num: i64=n; while bit!=0 { if num>=res+bit { num=num-(res+bit); res=(res/2)+bit } else { res=res/2 } bit=bit/4 } return res } called by 1: rmsnorm_fwd |
| 21 | func fp_exp(xq: i64) -> i64 { let y: i64=(xq*1512776)/S; var yi: i64=0; if y>=0 { yi=y/S } else { yi=0-(((0-y)+S-1)/S) } let yf: i64=y-yi*S; var p: i64=10085; p=58197+(p*yf)/S; p=251882+(p*yf)/S; p=726817+(p*yf)/S; p=S+(p*yf)/S; if yi>=0 { if yi>=31 { return 2000000000 } return p*(1<<yi) } let k: i64=0-yi; if k>=31 { return 0 } return p/(1<<k) } |
| 22 | func sigmoid(z: i64) -> i64 { let e: i64=fp_exp(0-z); return (S*S)/(S+e) } |
| 24 | func rmsnorm_fwd(x: *i64, gm: *i64, out: *i64, rms: *i64) -> i64 { var t: i64=0; while t<T { var ms: i64=0; var i: i64=0; while i<DM { ms=ms+x[t*DM+i]*x[t*DM+i]; i=i+1 } ms=ms/DM+EPS; let r: i64=isqrt(ms); rms[t]=r; let inv: i64=(S*S)/r; i=0; while i<DM { let nrm: i64=(x[t*DM+i]*inv)/S; out[t*DM+i]=(nrm*gm[i])/S; i=i+1 } t=t+1 } return 0 } |
| 25 | func rmsnorm_bwd(x: *i64, gm: *i64, dh: *i64, rms: *i64, dxacc: *i64, dgm: *i64) -> i64 { var t: i64=0; while t<T { let r: i64=rms[t]; let inv: i64=(S*S)/r; let invr3: i64=(((inv*inv)/S)*inv)/S; var c: i64=0; var i: i64=0; while i<DM { let nrm: i64=(x[t*DM+i]*inv)/S; dgm[i]=dgm[i]+(dh[t*DM+i]*nrm)/S; let dn: i64=(dh[t*DM+i]*gm[i])/S; c=c+(dn*x[t*DM+i])/S; i=i+1 } i=0; while i<DM { let dn: i64=(dh[t*DM+i]*gm[i])/S; let t1: i64=(dn*inv)/S; let tt: i64=(x[t*DM+i]*c)/S; let t2: i64=(((tt*invr3)/S))/DM; dxacc[t*DM+i]=dxacc[t*DM+i]+(t1-t2); i=i+1 } t=t+1 } return 0 } called by 1: layer_core_bwd |
| 26 | func attn_fwd(h: *i64, Wq: *i64, Wk: *i64, Wv: *i64, Q: *i64, K: *i64, V: *i64, A: *i64, O: *i64) -> i64 { var t: i64=0; while t<T { var i: i64=0; while i<DM { var aq: i64=0; var ak: i64=0; var av: i64=0; var k: i64=0; while k<DM { aq=aq+h[t*DM+k]*Wq[k*DM+i]; ak=ak+h[t*DM+k]*Wk[k*DM+i]; av=av+h[t*DM+k]*Wv[k*DM+i]; k=k+1 } Q[t*DM+i]=aq/S; K[t*DM+i]=ak/S; V[t*DM+i]=av/S; i=i+1 } t=t+1 } t=0; while t<T { var mx: i64=0-2000000000; var s: i64=0; while s<=t { var dot: i64=0; var i: i64=0; while i<DM { dot=dot+Q[t*DM+i]*K[s*DM+i]; i=i+1 } let sc: i64=((dot/S)*SCALE)/S; A[t*T+s]=sc; if sc>mx { mx=sc } s=s+1 } var sum: i64=0; s=0; while s<=t { let e: i64=fp_exp(A[t*T+s]-mx); A[t*T+s]=e; sum=sum+e; s=s+1 } s=0; while s<=t { A[t*T+s]=(A[t*T+s]*S+sum/2)/sum; s=s+1 } t=t+1 } t=0; while t<T { var i: i64=0; while i<DM { var acc: i64=0; var s: i64=0; while s<=t { acc=acc+A[t*T+s]*V[s*DM+i]; s=s+1 } O[t*DM+i]=acc/S; i=i+1 } t=t+1 } return 0 } |
| 27 | func attn_bwd(h: *i64, Wq: *i64, Wk: *i64, Wv: *i64, Q: *i64, K: *i64, V: *i64, A: *i64, dO: *i64, dWq: *i64, dWk: *i64, dWv: *i64, dh: *i64) -> i64 { let dV: *i64=sys_mmap(T*DM*8) as *i64; let dA: *i64=sys_mmap(T*T*8) as *i64; let dsc: *i64=sys_mmap(T*T*8) as *i64; let dQ: *i64=sys_mmap(T*DM*8) as *i64; let dK: *i64=sys_mmap(T*DM*8) as *i64; var s: i64=0; while s<DM*T { dV[s]=0; s=s+1 } s=0; while s<T { var i: i64=0; while i<DM { var acc: i64=0; var t: i64=s; while t<T { acc=acc+(A[t*T+s]*dO[t*DM+i])/S; t=t+1 } dV[s*DM+i]=acc; i=i+1 } s=s+1 } var t: i64=0; while t<T { s=0; while s<=t { var acc: i64=0; var i: i64=0; while i<DM { acc=acc+(dO[t*DM+i]*V[s*DM+i])/S; i=i+1 } dA[t*T+s]=acc; s=s+1 } t=t+1 } t=0; while t<T { var dot: i64=0; s=0; while s<=t { dot=dot+(A[t*T+s]*dA[t*T+s])/S; s=s+1 } s=0; while s<=t { dsc[t*T+s]=(A[t*T+s]*(dA[t*T+s]-dot))/S; s=s+1 } t=t+1 } t=0; while t<T*DM { dQ[t]=0; dK[t]=0; t=t+1 } t=0; while t<T { var i: i64=0; while i<DM { var acc: i64=0; s=0; while s<=t { let dqk: i64=(SCALE*dsc[t*T+s])/S; acc=acc+(dqk*K[s*DM+i])/S; s=s+1 } dQ[t*DM+i]=acc; i=i+1 } t=t+1 } s=0; while s<T { var i: i64=0; while i<DM { var acc: i64=0; t=s; while t<T { let dqk: i64=(SCALE*dsc[t*T+s])/S; acc=acc+(dqk*Q[t*DM+i])/S; t=t+1 } dK[s*DM+i]=acc; i=i+1 } s=s+1 } var kk: i64=0; while kk<DM { var i: i64=0; while i<DM { var aq: i64=0; var ak: i64=0; var av: i64=0; t=0; while t<T { aq=aq+(h[t*DM+kk]*dQ[t*DM+i])/S; ak=ak+(h[t*DM+kk]*dK[t*DM+i])/S; av=av+(h[t*DM+kk]*dV[t*DM+i])/S; t=t+1 } dWq[kk*DM+i]=aq; dWk[kk*DM+i]=ak; dWv[kk*DM+i]=av; i=i+1 } kk=kk+1 } t=0; while t<T { var k2: i64=0; while k2<DM { var acc: i64=0; var i: i64=0; while i<DM { acc=acc+(dQ[t*DM+i]*Wq[k2*DM+i])/S+(dK[t*DM+i]*Wk[k2*DM+i])/S+(dV[t*DM+i]*Wv[k2*DM+i])/S; i=i+1 } dh[t*DM+k2]=acc; k2=k2+1 } t=t+1 } return 0 } |
| 28 | func swiglu_fwd(h: *i64, Wg: *i64, Wu: *i64, Wd: *i64, g: *i64, a: *i64, u: *i64, m: *i64, ffn: *i64) -> i64 { var t: i64=0; while t<T { var j: i64=0; while j<HF { var gg: i64=0; var uu: i64=0; var k: i64=0; while k<DM { gg=gg+h[t*DM+k]*Wg[k*HF+j]; uu=uu+h[t*DM+k]*Wu[k*HF+j]; k=k+1 } let gv: i64=gg/S; g[t*HF+j]=gv; u[t*HF+j]=uu/S; let sig: i64=sigmoid(gv); let av: i64=(gv*sig)/S; a[t*HF+j]=av; m[t*HF+j]=(av*u[t*HF+j])/S; j=j+1 } t=t+1 } t=0; while t<T { var i: i64=0; while i<DM { var acc: i64=0; var j: i64=0; while j<HF { acc=acc+m[t*HF+j]*Wd[j*DM+i]; j=j+1 } ffn[t*DM+i]=acc/S; i=i+1 } t=t+1 } return 0 } |
| 29 | func swiglu_bwd(h: *i64, Wg: *i64, Wu: *i64, Wd: *i64, g: *i64, a: *i64, u: *i64, m: *i64, dOut: *i64, dWg: *i64, dWu: *i64, dWd: *i64, dh: *i64) -> i64 { let dm: *i64=sys_mmap(T*HF*8) as *i64; let dg: *i64=sys_mmap(T*HF*8) as *i64; let du: *i64=sys_mmap(T*HF*8) as *i64; var j: i64=0; while j<HF { var i: i64=0; while i<DM { var acc: i64=0; var t: i64=0; while t<T { acc=acc+(m[t*HF+j]*dOut[t*DM+i])/S; t=t+1 } dWd[j*DM+i]=acc; i=i+1 } j=j+1 } var t: i64=0; while t<T { j=0; while j<HF { var acc: i64=0; var i: i64=0; while i<DM { acc=acc+(dOut[t*DM+i]*Wd[j*DM+i])/S; i=i+1 } dm[t*HF+j]=acc; j=j+1 } t=t+1 } t=0; while t<T { j=0; while j<HF { let dmv: i64=dm[t*HF+j]; let dav: i64=(dmv*u[t*HF+j])/S; du[t*HF+j]=(dmv*a[t*HF+j])/S; let z: i64=g[t*HF+j]; let sig: i64=sigmoid(z); let zt: i64=(z*(S-sig))/S; let dsil: i64=(sig*(S+zt))/S; dg[t*HF+j]=(dav*dsil)/S; j=j+1 } t=t+1 } var k: i64=0; while k<DM { j=0; while j<HF { var ag: i64=0; var au: i64=0; t=0; while t<T { ag=ag+(h[t*DM+k]*dg[t*HF+j])/S; au=au+(h[t*DM+k]*du[t*HF+j])/S; t=t+1 } dWg[k*HF+j]=ag; dWu[k*HF+j]=au; j=j+1 } k=k+1 } t=0; while t<T { k=0; while k<DM { var acc: i64=0; j=0; while j<HF { acc=acc+(dg[t*HF+j]*Wg[k*HF+j])/S+(du[t*HF+j]*Wu[k*HF+j])/S; j=j+1 } dh[t*DM+k]=acc; k=k+1 } t=t+1 } return 0 } |
| 32 | func layer_core_fwd(P: *i64) -> i64 |
| 42 | func layer_core_bwd(P: *i64) -> i64 |
| 55 | func embed_fwd(P: *i64) -> i64 { let X: *i64=P[0] as *i64; let Em: *i64=P[9] as *i64; let tok: *i64=P[40] as *i64; var t: i64=0; while t<T { let id: i64=tok[t]; var i: i64=0; while i<DM { X[t*DM+i]=Em[id*DM+i]; i=i+1 } t=t+1 } return 0 } called by 1: main |
| 57 | func head_fwd(P: *i64) -> i64 |
| 70 | func head_bwd(P: *i64) -> i64 |
| 80 | func embed_bwd(P: *i64) -> i64 { let dEm: *i64=P[38] as *i64; let dX: *i64=P[34] as *i64; let tok: *i64=P[40] as *i64; var z: i64=0; while z<VOC*DM { dEm[z]=0; z=z+1 } var t: i64=0; while t<T { let id: i64=tok[t]; var i: i64=0; while i<DM { dEm[id*DM+i]=dEm[id*DM+i]+dX[t*DM+i]; i=i+1 } t=t+1 } return 0 } called by 1: main |
| 82 | func main() -> i64 |