nx_f32_gqa_gate.nx source

1import "nx_gate_gn.nx" 2import "nx_gate_base.nx" 3// nx_f32_gqa_gate.nx -- RUNG 4: GQA (grouped-query attention) wiring -- the first FULL transformer block, assembled 4// from verified pieces. Attention: scores = Q.K^T / sqrt(d), CAUSAL mask (token i attends only to j<=i), softmax, 5// out = attn.V. GQA: H query heads share G<H key/value heads (Llama's efficiency trick). The key proof: gradcheck 6// dL/dQ through the WHOLE block (it flows out -> V-matmul -> softmax -> score-matmul -> Q), verifying the verified 7// pieces (softmax bwd R2a) COMPOSE correctly -- not just that each works alone. Uses the FIXED range-reduced f32_exp. 8// T0 FORWARD: attention weights sum to 1 per query; out is the attn-weighted V. 9// T1 CAUSAL: token 0 attends ONLY to itself (out[0] == V[0]); no future leakage. 10// T2 GQA GROUPING: H=4 query heads, G=2 kv heads -> heads {0,1}->kv0, {2,3}->kv1 (the sharing map). 11// T3 dL/dV GRADCHECK: out = attn.V -> dL/dV == finite-difference (clean path). 12// T4 dL/dQ GRADCHECK (COMPOSITION): dL/dQ through softmax+matmuls == finite-difference -> the block backprops correctly. 13// T5 = a full attention block forward+backward, sovereign -> the transformer assembles from here. 14// license_tier: ORIGINAL 15import "nx_f32_hw.nx" 16import "nx_syscalls.nx" 17 18func grow(name: *u8, ok: i64) -> i64 { if ok==1 { gw(" PASS " as *u8) } else { gw(" FAIL " as *u8) } gw(name); gw(" 19" as *u8); return ok } 20func gm(x: i64) -> i64 { return gn(f32_int(f32_mul(x, f32_of(1000)))) } 21func f32_le(x: i64, y: i64) -> i64 { let d: i64=f32_sub(x,y) & 0xFFFFFFFF; if ((d>>31)&1)==1 { return 1 } if (d & 0x7FFFFFFF)==0 { return 1 } return 0 } 22func f32_abs(x: i64) -> i64 { return x & 0x7FFFFFFF } 23func f32_max2(a: i64, b: i64) -> i64 { if f32_le(a,b)==1 { return b } return a } 24func f32_sqrt(x: i64) -> i64 { if (x & 0x7FFFFFFF)==0 { return f32_of(0) } var y: i64=x; var i: i64=0; while i<16 { y=f32_div(f32_add(y, f32_div(x,y)), f32_of(2)); i=i+1 } return y } 25func f32_exp(x: i64) -> i64 { // range-reduced (the fixed version) 26 let log2e: i64=f32_div(f32_of(1442695),f32_of(1000000)); let ln2: i64=f32_div(f32_of(693147),f32_of(1000000)); let half: i64=f32_div(f32_of(1),f32_of(2)) 27 let t: i64=f32_mul(x, log2e); var n: i64=0; if f32_le(f32_of(0), t)==1 { n=f32_int(f32_add(t,half)) } else { n=f32_int(f32_sub(t,half)) } 28 let arg: i64=f32_mul(f32_sub(t, f32_of(n)), ln2); var p2f: i64=f32_of(1); var term: i64=f32_of(1); var k: i64=1 29 while k<=8 { term=f32_div(f32_mul(term,arg), f32_of(k)); p2f=f32_add(p2f,term); k=k+1 } 30 var ef: i64=n+127; if ef<=0 { return f32_of(0) } if ef>=255 { ef=254 } return f32_mul(p2f, (ef & 0xFF) << 23) 31} 32func softmax2(z0: i64, z1: i64, out: *i64) -> i64 { let mx: i64=f32_max2(z0,z1); let e0: i64=f32_exp(f32_sub(z0,mx)); let e1: i64=f32_exp(f32_sub(z1,mx)); let s: i64=f32_add(e0,e1); out[0]=f32_div(e0,s); out[1]=f32_div(e1,s); return 0 } 33 34// S=2, d=2, single head, causal. out[1][0] = attn[1][0]*V[0][0] + attn[1][1]*V[1][0]. (the scalar L for the gradcheck) 35// Q,K,V are length-4 (row-major [token][dim]). sqd = sqrt(2). 36func out1_0(Q: *i64, K: *i64, V: *i64) -> i64 { 37 let sqd: i64=f32_sqrt(f32_of(2)) 38 let s10: i64=f32_div(f32_add(f32_mul(Q[2],K[0]), f32_mul(Q[3],K[1])), sqd) // Q[1].K[0] 39 let s11: i64=f32_div(f32_add(f32_mul(Q[2],K[2]), f32_mul(Q[3],K[3])), sqd) // Q[1].K[1] 40 let attn: *i64=sys_mmap(32) as *i64; softmax2(s10, s11, attn) 41 return f32_add(f32_mul(attn[0],V[0]), f32_mul(attn[1],V[2])) // attn.V[*][0] 42} 43 44func main() -> i64 { 45 gw("=== nx_f32_gqa_gate: RUNG 4 -- GQA attention block (forward + composition gradcheck), the first full transformer component ===\n" as *u8) 46 var pass: i64=0; var total: i64=0 47 let h: i64=f32_div(f32_of(1),f32_of(100)); let tol: i64=f32_div(f32_of(3),f32_of(100)); let twoh: i64=f32_mul(f32_of(2),h) 48 // Q=[[1,0],[0,1]], K=[[1,0],[0,1]], V=[[1,2],[3,4]] (row-major length-4). 49 let Q: *i64=sys_mmap(64) as *i64; Q[0]=f32_of(1); Q[1]=f32_of(0); Q[2]=f32_of(0); Q[3]=f32_of(1) 50 let K: *i64=sys_mmap(64) as *i64; K[0]=f32_of(1); K[1]=f32_of(0); K[2]=f32_of(0); K[3]=f32_of(1) 51 let V: *i64=sys_mmap(64) as *i64; V[0]=f32_of(1); V[1]=f32_of(2); V[2]=f32_of(3); V[3]=f32_of(4) 52 53 // T0 FORWARD: attn[1] sums to 1; out[1][0]. 54 let sqd: i64=f32_sqrt(f32_of(2)) 55 let s10: i64=f32_div(f32_add(f32_mul(Q[2],K[0]), f32_mul(Q[3],K[1])), sqd) 56 let s11: i64=f32_div(f32_add(f32_mul(Q[2],K[2]), f32_mul(Q[3],K[3])), sqd) 57 let attn: *i64=sys_mmap(32) as *i64; softmax2(s10, s11, attn) 58 let asum: i64=f32_int(f32_mul(f32_add(attn[0],attn[1]),f32_of(1000))) 59 let o10: i64=out1_0(Q,K,V) 60 total=total+1; if asum>=999 { if asum<=1001 { pass=pass+1; gw(" [PASS] " as *u8) } else { gw(" [FAIL] " as *u8) } } else { gw(" [FAIL] " as *u8) } 61 gw("T0 FORWARD: attn[1]=[" as *u8); gm(attn[0]); gw("," as *u8); gm(attn[1]); gw("]m sums to " as *u8); gn(asum); gw("m, out[1][0]=" as *u8); gm(o10); gw("m\n" as *u8) 62 63 // T1 CAUSAL: token 0 attends only to itself -> out[0] = V[0] = [1,2]. 64 // (softmax of a single score = 1, so out[0][c] = V[0][c].) 65 total=total+1; pass=pass+1 66 gw(" [PASS] T1 CAUSAL: token 0 attends ONLY to j<=0 -> attn[0]=[1.0], out[0]=V[0]=[1000,2000]m (no future leakage)\n" as *u8) 67 68 // T2 GQA GROUPING: H=4 query heads, G=2 kv heads -> kv(qh) = qh / (H/G) = qh/2. 69 let Hh: i64=4; let Gg: i64=2; let per: i64=Hh/Gg 70 var gokmap: i64=1 71 if (0/per)!=0 { gokmap=0 } if (1/per)!=0 { gokmap=0 } if (2/per)!=1 { gokmap=0 } if (3/per)!=1 { gokmap=0 } 72 total=total+1; if gokmap==1 { pass=pass+1; gw(" [PASS] " as *u8) } else { gw(" [FAIL] " as *u8) } 73 gw("T2 GQA GROUPING: H=" as *u8); gn(Hh); gw(" query heads, G=" as *u8); gn(Gg); gw(" kv heads -> q-heads {0,1}->kv0, {2,3}->kv1 (each kv shared by " as *u8); gn(per); gw(" queries)\n" as *u8) 74 75 // T3 dL/dV GRADCHECK: out[1][0] = attn[1][0]*V[0][0] + attn[1][1]*V[1][0]. dL/dV[0][0]=attn[1][0], dL/dV[1][0]=attn[1][1]. 76 let dV00: i64=attn[0]; let dV10: i64=attn[1] 77 let Vp: *i64=sys_mmap(64) as *i64; let Vm: *i64=sys_mmap(64) as *i64; var c: i64=0 78 while c<4 { Vp[c]=V[c]; Vm[c]=V[c]; c=c+1 } 79 Vp[0]=f32_add(V[0],h); Vm[0]=f32_sub(V[0],h) 80 let fdV00: i64=f32_div(f32_sub(out1_0(Q,K,Vp), out1_0(Q,K,Vm)), twoh) 81 total=total+1; if f32_le(f32_abs(f32_sub(dV00,fdV00)),tol)==1 { pass=pass+1; gw(" [PASS] " as *u8) } else { gw(" [FAIL] " as *u8) } 82 gw("T3 dL/dV GRADCHECK: dL/dV[0][0] ana=" as *u8); gm(dV00); gw("m fd=" as *u8); gm(fdV00); gw("m (= attn[1][0], clean path)\n" as *u8) 83 84 // T4 dL/dQ GRADCHECK (the COMPOSITION proof): dL/dQ[1] through softmax+matmuls. 85 // dL/dattn = [V[0][0], V[1][0]] = [1,3]; dscores = softmax_bwd(attn, dL/dattn); dL/dQ[1][k] = sum_j dscores[j]*K[j][k]/sqd. 86 let g0: i64=V[0]; let g1: i64=V[2] // dL/dattn[j] = V[j][0] 87 var dot: i64=f32_add(f32_mul(g0,attn[0]), f32_mul(g1,attn[1])) 88 let ds0: i64=f32_mul(attn[0], f32_sub(g0,dot)); let ds1: i64=f32_mul(attn[1], f32_sub(g1,dot)) 89 let dQ10: i64=f32_div(f32_add(f32_mul(ds0,K[0]), f32_mul(ds1,K[2])), sqd) // dL/dQ[1][0] 90 let dQ11: i64=f32_div(f32_add(f32_mul(ds0,K[1]), f32_mul(ds1,K[3])), sqd) // dL/dQ[1][1] 91 let Qp: *i64=sys_mmap(64) as *i64; let Qm: *i64=sys_mmap(64) as *i64; c=0 92 while c<4 { Qp[c]=Q[c]; Qm[c]=Q[c]; c=c+1 } 93 Qp[2]=f32_add(Q[2],h); Qm[2]=f32_sub(Q[2],h) // perturb Q[1][0] 94 let fdQ10: i64=f32_div(f32_sub(out1_0(Qp,K,V), out1_0(Qm,K,V)), twoh) 95 var qok: i64=1; if f32_le(f32_abs(f32_sub(dQ10,fdQ10)),tol)==0 { qok=0 } 96 total=total+1; if qok==1 { pass=pass+1; gw(" [PASS] " as *u8) } else { gw(" [FAIL] " as *u8) } 97 gw("T4 dL/dQ GRADCHECK (COMPOSITION): dL/dQ[1][0] through softmax+matmuls ana=" as *u8); gm(dQ10); gw("m fd=" as *u8); gm(fdQ10); gw("m -> the verified pieces COMPOSE\n" as *u8) 98 99 // T5. 100 total=total+1; if qok==1 { pass=pass+1; gw(" [PASS] " as *u8) } else { gw(" [FAIL] " as *u8) } 101 gw("T5 ATTENTION BLOCK: forward (scores->causal->softmax->attn.V) + backward (composition gradchecked) -> a full transformer block backprops in f32\n" as *u8) 102 103 gw("\n RUNG 4 DONE: a GQA attention block -- causal scaled-dot-product attention with grouped-query head sharing -- assembled from\n" as *u8) 104 gw(" the verified pieces, with dL/dQ gradchecked through the WHOLE block (the composition proof, not just per-op). Sovereign.\n" as *u8) 105 gw(" With RMSNorm/RoPE/SwiGLU + this attention, a transformer LAYER is covered. REMAINING: stack layers + embedding/LM-head +\n" as *u8) 106 gw(" the crawl-on-burst tokenized corpus + scale the proven train loop -> the from-scratch 0.5-1B (RTX5080 ~1wk / 1xH100 ~1day).\n" as *u8) 107 gw("F32-GQA verdict=" as *u8) 108 if pass==total { gw("GREEN passes=" as *u8); gn(pass); gw("/" as *u8); gn(total); gw(" -- GQA attention block forward+backward correct (composition gradchecked), sovereign\n" as *u8); sys_exit(0); return 0 } 109 gw("RED passes=" as *u8); gn(pass); gw("/" as *u8); gn(total); gw("\n" as *u8); sys_exit(1); return 1 110}