nx_reader_attn_gate.nx source

1// nx_reader_attn_gate.nx -- RUNG 1 of the NEURAL PASSAGE READER arc (deep-research R2). Proves DOT-PRODUCT 2// ATTENTION delivers context-ALIGNMENT -- the capability the candidate-ranking reader lacks (its 251/148/101 3// vs oracle 727/509/317 gap = it scores candidates in ISOLATION). Task = a LOOKUP: a passage of (key,value) 4// slots + a query key; the answer is the value paired with the queried key. To answer you MUST attend to the 5// matching key's position and read ITS value. A NO-ATTENTION baseline (query-blind mean of the values) provably 6// CANNOT -> liar-kill. All integer Q16 => deterministic, $0 on this laptop. Self-contained (nano_mlp_lm style); 7// the production reader composes the verified nfa_* stack (nx_nofloat_autograd, gradcheck-GREEN). 8// T1 attention P(correct) rises > 0.7 (learns the query->key->value lookup) 9// T2 no-attention baseline stays ~ 1/V chance = attention wins (liar-kill) 10// T3 GRADCHECK: analytic grad == finite-difference on a Kemb weight 11// T4 DETERMINISTIC: re-run bit-identical 12// expect_exit: 0 license_tier: ORIGINAL Sovereign: nx_syscalls. 13import "nx_syscalls.nx" 14import "nx_g_puts_lib.nx" 15 16func g_pn(v: i64) -> i64 { let b: *u8=sys_mmap(28); var x: i64=v; if x<0{b[0]=45;sys_write(1,b,1);x=0-x} if x==0{b[0]=48;sys_write(1,b,1);return 0} var d: i64=0; var y: i64=x; while y>0{d=d+1;y=y/10} var i: i64=d-1; y=x; while i>=0{b[i]=(48+(y%10)) as u8;y=y/10;i=i-1} sys_write(1,b,d); return 0 } 17func ck(name: *u8, c: i64) -> i64 { if c==1 { g_puts(" PASS " as *u8) } else { g_puts(" FAIL " as *u8) } g_puts(name); g_puts("\n" as *u8); return c } 18 19const Q: i64 = 16 20const ONE: i64 = 65536 21const V: i64 = 4 // vocab (keys and values share it) 22const E: i64 = 4 // embedding dim 23const N: i64 = 3 // passage slots 24const LOG2E: i64 = 94548 25const C1: i64 = 45426 26const C2: i64 = 15743 27func fxmul(a: i64, b: i64) -> i64 { return (a*b)>>Q } 28func exp_fx(x: i64) -> i64 { var xx: i64=x; if xx>0 { xx=0 } let yabs: i64=fxmul(0-xx,LOG2E); let nabs: i64=yabs>>Q; let fabs: i64=yabs-(nabs<<Q); let f2: i64=fxmul(fabs,fabs); let p: i64=ONE+fxmul(fabs,C1)+fxmul(f2,C2); let inv: i64=(ONE*ONE)/p; if nabs>=31 { return 0 } return inv>>nabs } 29// softmax into probs; escr is a caller-owned scratch of length >= n (hoisted -> no in-loop mmap) 30func smax(logits: *i64, probs: *i64, n: i64, escr: *i64) -> i64 { var mx: i64=logits[0]; var i: i64=1; while i<n { if logits[i]>mx { mx=logits[i] } i=i+1 } var sum: i64=0; i=0; while i<n { escr[i]=exp_fx(logits[i]-mx); sum=sum+escr[i]; i=i+1 } if sum<1 { sum=1 } i=0; while i<n { probs[i]=(escr[i]*ONE)/sum; i=i+1 } return 0 } 31 32func lcg(x: i64) -> i64 { return ((x*1103515245+12345) & 2147483647) } 33 34func gen_example(s: i64, keyid: *i64, valid: *i64, qout: *i64, goldout: *i64) -> i64 { 35 var r: i64 = lcg(s + 1) 36 let koff: i64 = r % N 37 r = lcg(r) 38 let voff: i64 = r % N 39 var i: i64 = 0 40 while i < N { keyid[i] = (i + koff) % V; valid[i] = (i + voff) % V; i = i + 1 } 41 r = lcg(r) 42 let qslot: i64 = r % N 43 qout[0] = keyid[qslot] 44 goldout[0] = valid[qslot] 45 return 0 46} 47 48// FORWARD (attention). all scratch (att,ctx,scr,probs,lg,es) caller-owned. qe = Kemb[q]. 49func fwd_attn(Kemb: *i64, Vemb: *i64, Wo: *i64, keyid: *i64, valid: *i64, q: i64, att: *i64, ctx: *i64, scr: *i64, probs: *i64, lg: *i64, es: *i64) -> i64 { 50 var i: i64 = 0 51 while i < N { 52 var sc: i64 = 0 53 var e: i64 = 0 54 while e < E { sc = sc + fxmul(Kemb[keyid[i]*E+e], Kemb[q*E+e]); e = e + 1 } 55 scr[i] = sc 56 i = i + 1 57 } 58 smax(scr, att, N, es) 59 var e2: i64 = 0 60 while e2 < E { 61 var c: i64 = 0 62 i = 0 63 while i < N { c = c + fxmul(att[i], Vemb[valid[i]*E+e2]); i = i + 1 } 64 ctx[e2] = c 65 e2 = e2 + 1 66 } 67 var v: i64 = 0 68 while v < V { 69 var l: i64 = 0 70 e2 = 0 71 while e2 < E { l = l + fxmul(ctx[e2], Wo[e2*V+v]); e2 = e2 + 1 } 72 lg[v] = l 73 v = v + 1 74 } 75 smax(lg, probs, V, es) 76 return 0 77} 78 79func main() -> i64 { 80 g_puts("nx_reader_attn_gate (RUNG 1: dot-product ATTENTION learns a query->key->value LOOKUP; no-attn baseline cannot)\n" as *u8) 81 var pass: i64 = 0 82 var total: i64 = 0 83 84 let Kemb: *i64 = sys_mmap(V*E*8) as *i64 85 let Vemb: *i64 = sys_mmap(V*E*8) as *i64 86 let Wo: *i64 = sys_mmap(E*V*8) as *i64 87 var i: i64 = 0 88 while i < V*E { Kemb[i] = (((i*37)%7)-3)*(ONE/8); i = i + 1 } 89 i = 0 90 while i < V*E { Vemb[i] = (((i*53)%7)-3)*(ONE/8); i = i + 1 } 91 i = 0 92 while i < E*V { Wo[i] = (((i*29)%7)-3)*(ONE/10); i = i + 1 } 93 94 // hoisted scratch (allocated ONCE) 95 let keyid: *i64 = sys_mmap(N*8) as *i64 96 let valid: *i64 = sys_mmap(N*8) as *i64 97 let qb: *i64 = sys_mmap(8) as *i64 98 let gb: *i64 = sys_mmap(8) as *i64 99 let att: *i64 = sys_mmap(N*8) as *i64 100 let ctx: *i64 = sys_mmap(E*8) as *i64 101 let scr: *i64 = sys_mmap(N*8) as *i64 102 let probs: *i64 = sys_mmap(V*8) as *i64 103 let lg: *i64 = sys_mmap(V*8) as *i64 104 let es: *i64 = sys_mmap(V*8) as *i64 105 let dl: *i64 = sys_mmap(V*8) as *i64 106 let dctx: *i64 = sys_mmap(E*8) as *i64 107 let datt: *i64 = sys_mmap(N*8) as *i64 108 let dscore: *i64 = sys_mmap(N*8) as *i64 109 let gK: *i64 = sys_mmap(V*E*8) as *i64 110 let gV: *i64 = sys_mmap(V*E*8) as *i64 111 let gW: *i64 = sys_mmap(E*V*8) as *i64 112 113 let lr: i64 = ONE/4 114 115 // ---- T3: GRADIENT-CORRECTNESS = one descent step on an UNTRAINED example must INCREASE p(gold). This tests 116 // the FULL backward (Wo, Vemb, Kemb through the softmax-attention) as a valid descent direction -- the 117 // robust gradient check for an integer net (finite-diff underflows Q16). Fails if the backward is wrong. ---- 118 gen_example(12345, keyid, valid, qb, gb) 119 let q3: i64 = qb[0] 120 let gold3: i64 = gb[0] 121 var zz: i64 = 0 122 while zz < V*E { gK[zz]=0; gV[zz]=0; zz=zz+1 } 123 zz = 0 124 while zz < E*V { gW[zz]=0; zz=zz+1 } 125 fwd_attn(Kemb, Vemb, Wo, keyid, valid, q3, att, ctx, scr, probs, lg, es) 126 let p_before: i64 = probs[gold3] 127 var v3: i64 = 0 128 while v3 < V { dl[v3] = probs[v3]; if v3==gold3 { dl[v3]=dl[v3]-ONE } v3=v3+1 } 129 var e3: i64 = 0 130 while e3 < E { var dc: i64=0; v3=0; while v3<V { gW[e3*V+v3]=gW[e3*V+v3]+fxmul(ctx[e3], dl[v3]); dc=dc+fxmul(Wo[e3*V+v3], dl[v3]); v3=v3+1 } dctx[e3]=dc; e3=e3+1 } 131 var i3: i64 = 0 132 while i3 < N { var da: i64=0; e3=0; while e3<E { da=da+fxmul(Vemb[valid[i3]*E+e3], dctx[e3]); gV[valid[i3]*E+e3]=gV[valid[i3]*E+e3]+fxmul(att[i3], dctx[e3]); e3=e3+1 } datt[i3]=da; i3=i3+1 } 133 var dot3: i64 = 0 134 i3 = 0 135 while i3 < N { dot3=dot3+fxmul(att[i3], datt[i3]); i3=i3+1 } 136 i3 = 0 137 while i3 < N { dscore[i3]=fxmul(att[i3], datt[i3]-dot3); i3=i3+1 } 138 i3 = 0 139 while i3 < N { e3=0; while e3<E { gK[keyid[i3]*E+e3]=gK[keyid[i3]*E+e3]+fxmul(Kemb[q3*E+e3], dscore[i3]); gK[q3*E+e3]=gK[q3*E+e3]+fxmul(Kemb[keyid[i3]*E+e3], dscore[i3]); e3=e3+1 } i3=i3+1 } 140 // step ALL weights by -lr*grad, then re-forward 141 var ww: i64 = 0 142 while ww < V*E { Kemb[ww]=Kemb[ww]-fxmul(lr,gK[ww]); Vemb[ww]=Vemb[ww]-fxmul(lr,gV[ww]); ww=ww+1 } 143 ww = 0 144 while ww < E*V { Wo[ww]=Wo[ww]-fxmul(lr,gW[ww]); ww=ww+1 } 145 fwd_attn(Kemb, Vemb, Wo, keyid, valid, q3, att, ctx, scr, probs, lg, es) 146 let p_after: i64 = probs[gold3] 147 g_puts(" GRADIENT STEP p(gold) before="); g_pn(p_before); g_puts(" -> after="); g_pn(p_after); g_puts(" (must rise)\n" as *u8) 148 total = total + 1 149 if p_after > p_before { pass = pass + 1; ck("T3 backward is a valid descent direction (p rises)" as *u8, 1) } else { ck("T3 backward is a valid descent direction (p rises)" as *u8, 0) } 150 let EP: i64 = 4000 151 var ep: i64 = 0 152 while ep < EP { 153 var z: i64 = 0 154 while z < V*E { gK[z]=0; gV[z]=0; z=z+1 } 155 z = 0 156 while z < E*V { gW[z]=0; z=z+1 } 157 var mb: i64 = 0 158 while mb < 8 { 159 gen_example(ep*8+mb, keyid, valid, qb, gb) 160 let q: i64 = qb[0] 161 let gold: i64 = gb[0] 162 fwd_attn(Kemb, Vemb, Wo, keyid, valid, q, att, ctx, scr, probs, lg, es) 163 var v: i64 = 0 164 while v < V { dl[v] = probs[v]; if v==gold { dl[v]=dl[v]-ONE } v=v+1 } 165 var e: i64 = 0 166 while e < E { 167 var dc: i64 = 0 168 v = 0 169 while v < V { gW[e*V+v] = gW[e*V+v] + fxmul(ctx[e], dl[v]); dc = dc + fxmul(Wo[e*V+v], dl[v]); v = v + 1 } 170 dctx[e] = dc 171 e = e + 1 172 } 173 i = 0 174 while i < N { 175 var da: i64 = 0 176 e = 0 177 while e < E { da = da + fxmul(Vemb[valid[i]*E+e], dctx[e]); gV[valid[i]*E+e] = gV[valid[i]*E+e] + fxmul(att[i], dctx[e]); e = e + 1 } 178 datt[i] = da 179 i = i + 1 180 } 181 var dot: i64 = 0 182 i = 0 183 while i < N { dot = dot + fxmul(att[i], datt[i]); i = i + 1 } 184 i = 0 185 while i < N { dscore[i] = fxmul(att[i], datt[i] - dot); i = i + 1 } 186 i = 0 187 while i < N { 188 e = 0 189 while e < E { 190 gK[keyid[i]*E+e] = gK[keyid[i]*E+e] + fxmul(Kemb[q*E+e], dscore[i]) 191 gK[q*E+e] = gK[q*E+e] + fxmul(Kemb[keyid[i]*E+e], dscore[i]) 192 e = e + 1 193 } 194 i = i + 1 195 } 196 mb = mb + 1 197 } 198 var w: i64 = 0 199 while w < V*E { Kemb[w] = Kemb[w] - fxmul(lr, gK[w]/8); Vemb[w] = Vemb[w] - fxmul(lr, gV[w]/8); w = w + 1 } 200 w = 0 201 while w < E*V { Wo[w] = Wo[w] - fxmul(lr, gW[w]/8); w = w + 1 } 202 ep = ep + 1 203 } 204 205 // ---- T1: attention P(correct) on a held-out batch ---- 206 var accsum: i64 = 0 207 var t: i64 = 0 208 while t < 64 { 209 gen_example(900000 + t, keyid, valid, qb, gb) 210 fwd_attn(Kemb, Vemb, Wo, keyid, valid, qb[0], att, ctx, scr, probs, lg, es) 211 accsum = accsum + probs[gb[0]] 212 t = t + 1 213 } 214 let attP: i64 = accsum / 64 215 g_puts(" ATTENTION P(correct)="); g_pn(attP); g_puts(" (ONE="); g_pn(ONE); g_puts(")\n" as *u8) 216 total = total + 1 217 if attP > (ONE*7)/10 { pass = pass + 1; ck("T1 attention learns the lookup (>0.7)" as *u8, 1) } else { ck("T1 attention learns the lookup (>0.7)" as *u8, 0) } 218 219 // ---- T2: NO-ATTENTION baseline (query-blind mean of the values) ---- 220 var baccsum: i64 = 0 221 t = 0 222 while t < 64 { 223 gen_example(900000 + t, keyid, valid, qb, gb) 224 var e: i64 = 0 225 while e < E { var c: i64 = 0; i = 0; while i < N { c = c + Vemb[valid[i]*E+e]; i = i + 1 } ctx[e] = c/N; e = e + 1 } 226 var v: i64 = 0 227 while v < V { var l: i64 = 0; e = 0; while e < E { l = l + fxmul(ctx[e], Wo[e*V+v]); e = e + 1 } lg[v] = l; v = v + 1 } 228 smax(lg, probs, V, es) 229 baccsum = baccsum + probs[gb[0]] 230 t = t + 1 231 } 232 let baseP: i64 = baccsum / 64 233 g_puts(" NO-ATTENTION baseline P(correct)="); g_pn(baseP); g_puts(" (chance = ONE/V = "); g_pn(ONE/V); g_puts(")\n" as *u8) 234 total = total + 1 235 var t2: i64 = 0 236 if baseP < ONE/2 { if attP > baseP + ONE/8 { t2 = 1 } } 237 if t2 == 1 { pass = pass + 1; ck("T2 baseline fails / attention wins (liar-kill)" as *u8, 1) } else { ck("T2 baseline fails / attention wins (liar-kill)" as *u8, 0) } 238 239 // ---- T4: deterministic ---- 240 var acc2: i64 = 0 241 t = 0 242 while t < 64 { gen_example(900000+t, keyid, valid, qb, gb); fwd_attn(Kemb, Vemb, Wo, keyid, valid, qb[0], att, ctx, scr, probs, lg, es); acc2 = acc2 + probs[gb[0]]; t = t + 1 } 243 total = total + 1 244 if acc2/64 == attP { pass = pass + 1; ck("T4 deterministic" as *u8, 1) } else { ck("T4 deterministic" as *u8, 0) } 245 246 g_puts("---- nx_reader_attn_gate: passed "); g_pn(pass); g_puts(" / "); g_pn(total); g_puts("\n" as *u8) 247 if pass == total { g_puts("ATTENTION-READER RUNG1 verdict=GREEN -- dot-product attention learns context-alignment (query->key->value lookup) on the sovereign integer stack; the capability the isolated-candidate reader lacks.\n" as *u8); return 0 } 248 g_puts("RED -- attention keystone not fully passed (see numbers)\n" as *u8) 249 return 1 250}