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" 15import "nx_gate_verdict.nx" 16 17func 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 } 18func 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 } 19 20const Q: i64 = 16 21const ONE: i64 = 65536 22const V: i64 = 4 // vocab (keys and values share it) 23const E: i64 = 4 // embedding dim 24const N: i64 = 3 // passage slots 25const LOG2E: i64 = 94548 26const C1: i64 = 45426 27const C2: i64 = 15743 28func fxmul(a: i64, b: i64) -> i64 { return (a*b)>>Q } 29func 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 } 30// softmax into probs; escr is a caller-owned scratch of length >= n (hoisted -> no in-loop mmap) 31func 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 } 32 33func lcg(x: i64) -> i64 { return ((x*1103515245+12345) & 2147483647) } 34 35func gen_example(s: i64, keyid: *i64, valid: *i64, qout: *i64, goldout: *i64) -> i64 { 36 var r: i64 = lcg(s + 1) 37 let koff: i64 = r % N 38 r = lcg(r) 39 let voff: i64 = r % N 40 var i: i64 = 0 41 while i < N { keyid[i] = (i + koff) % V; valid[i] = (i + voff) % V; i = i + 1 } 42 r = lcg(r) 43 let qslot: i64 = r % N 44 qout[0] = keyid[qslot] 45 goldout[0] = valid[qslot] 46 return 0 47} 48 49// FORWARD (attention). all scratch (att,ctx,scr,probs,lg,es) caller-owned. qe = Kemb[q]. 50func 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 { 51 var i: i64 = 0 52 while i < N { 53 var sc: i64 = 0 54 var e: i64 = 0 55 while e < E { sc = sc + fxmul(Kemb[keyid[i]*E+e], Kemb[q*E+e]); e = e + 1 } 56 scr[i] = sc 57 i = i + 1 58 } 59 smax(scr, att, N, es) 60 var e2: i64 = 0 61 while e2 < E { 62 var c: i64 = 0 63 i = 0 64 while i < N { c = c + fxmul(att[i], Vemb[valid[i]*E+e2]); i = i + 1 } 65 ctx[e2] = c 66 e2 = e2 + 1 67 } 68 var v: i64 = 0 69 while v < V { 70 var l: i64 = 0 71 e2 = 0 72 while e2 < E { l = l + fxmul(ctx[e2], Wo[e2*V+v]); e2 = e2 + 1 } 73 lg[v] = l 74 v = v + 1 75 } 76 smax(lg, probs, V, es) 77 return 0 78} 79 80func main() -> i64 { 81 g_puts("nx_reader_attn_gate (RUNG 1: dot-product ATTENTION learns a query->key->value LOOKUP; no-attn baseline cannot)\n" as *u8) 82 var pass: i64 = 0 83 var total: i64 = 0 84 85 let Kemb: *i64 = sys_mmap(V*E*8) as *i64 86 let Vemb: *i64 = sys_mmap(V*E*8) as *i64 87 let Wo: *i64 = sys_mmap(E*V*8) as *i64 88 var i: i64 = 0 89 while i < V*E { Kemb[i] = (((i*37)%7)-3)*(ONE/8); i = i + 1 } 90 i = 0 91 while i < V*E { Vemb[i] = (((i*53)%7)-3)*(ONE/8); i = i + 1 } 92 i = 0 93 while i < E*V { Wo[i] = (((i*29)%7)-3)*(ONE/10); i = i + 1 } 94 95 // hoisted scratch (allocated ONCE) 96 let keyid: *i64 = sys_mmap(N*8) as *i64 97 let valid: *i64 = sys_mmap(N*8) as *i64 98 let qb: *i64 = sys_mmap(8) as *i64 99 let gb: *i64 = sys_mmap(8) as *i64 100 let att: *i64 = sys_mmap(N*8) as *i64 101 let ctx: *i64 = sys_mmap(E*8) as *i64 102 let scr: *i64 = sys_mmap(N*8) as *i64 103 let probs: *i64 = sys_mmap(V*8) as *i64 104 let lg: *i64 = sys_mmap(V*8) as *i64 105 let es: *i64 = sys_mmap(V*8) as *i64 106 let dl: *i64 = sys_mmap(V*8) as *i64 107 let dctx: *i64 = sys_mmap(E*8) as *i64 108 let datt: *i64 = sys_mmap(N*8) as *i64 109 let dscore: *i64 = sys_mmap(N*8) as *i64 110 let gK: *i64 = sys_mmap(V*E*8) as *i64 111 let gV: *i64 = sys_mmap(V*E*8) as *i64 112 let gW: *i64 = sys_mmap(E*V*8) as *i64 113 114 let lr: i64 = ONE/4 115 116 // ---- T3: GRADIENT-CORRECTNESS = one descent step on an UNTRAINED example must INCREASE p(gold). This tests 117 // the FULL backward (Wo, Vemb, Kemb through the softmax-attention) as a valid descent direction -- the 118 // robust gradient check for an integer net (finite-diff underflows Q16). Fails if the backward is wrong. ---- 119 gen_example(12345, keyid, valid, qb, gb) 120 let q3: i64 = qb[0] 121 let gold3: i64 = gb[0] 122 var zz: i64 = 0 123 while zz < V*E { gK[zz]=0; gV[zz]=0; zz=zz+1 } 124 zz = 0 125 while zz < E*V { gW[zz]=0; zz=zz+1 } 126 fwd_attn(Kemb, Vemb, Wo, keyid, valid, q3, att, ctx, scr, probs, lg, es) 127 let p_before: i64 = probs[gold3] 128 var v3: i64 = 0 129 while v3 < V { dl[v3] = probs[v3]; if v3==gold3 { dl[v3]=dl[v3]-ONE } v3=v3+1 } 130 var e3: i64 = 0 131 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 } 132 var i3: i64 = 0 133 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 } 134 var dot3: i64 = 0 135 i3 = 0 136 while i3 < N { dot3=dot3+fxmul(att[i3], datt[i3]); i3=i3+1 } 137 i3 = 0 138 while i3 < N { dscore[i3]=fxmul(att[i3], datt[i3]-dot3); i3=i3+1 } 139 i3 = 0 140 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 } 141 // step ALL weights by -lr*grad, then re-forward 142 var ww: i64 = 0 143 while ww < V*E { Kemb[ww]=Kemb[ww]-fxmul(lr,gK[ww]); Vemb[ww]=Vemb[ww]-fxmul(lr,gV[ww]); ww=ww+1 } 144 ww = 0 145 while ww < E*V { Wo[ww]=Wo[ww]-fxmul(lr,gW[ww]); ww=ww+1 } 146 fwd_attn(Kemb, Vemb, Wo, keyid, valid, q3, att, ctx, scr, probs, lg, es) 147 let p_after: i64 = probs[gold3] 148 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) 149 total = total + 1 150 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) } 151 let EP: i64 = 4000 152 var ep: i64 = 0 153 while ep < EP { 154 var z: i64 = 0 155 while z < V*E { gK[z]=0; gV[z]=0; z=z+1 } 156 z = 0 157 while z < E*V { gW[z]=0; z=z+1 } 158 var mb: i64 = 0 159 while mb < 8 { 160 gen_example(ep*8+mb, keyid, valid, qb, gb) 161 let q: i64 = qb[0] 162 let gold: i64 = gb[0] 163 fwd_attn(Kemb, Vemb, Wo, keyid, valid, q, att, ctx, scr, probs, lg, es) 164 var v: i64 = 0 165 while v < V { dl[v] = probs[v]; if v==gold { dl[v]=dl[v]-ONE } v=v+1 } 166 var e: i64 = 0 167 while e < E { 168 var dc: i64 = 0 169 v = 0 170 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 } 171 dctx[e] = dc 172 e = e + 1 173 } 174 i = 0 175 while i < N { 176 var da: i64 = 0 177 e = 0 178 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 } 179 datt[i] = da 180 i = i + 1 181 } 182 var dot: i64 = 0 183 i = 0 184 while i < N { dot = dot + fxmul(att[i], datt[i]); i = i + 1 } 185 i = 0 186 while i < N { dscore[i] = fxmul(att[i], datt[i] - dot); i = i + 1 } 187 i = 0 188 while i < N { 189 e = 0 190 while e < E { 191 gK[keyid[i]*E+e] = gK[keyid[i]*E+e] + fxmul(Kemb[q*E+e], dscore[i]) 192 gK[q*E+e] = gK[q*E+e] + fxmul(Kemb[keyid[i]*E+e], dscore[i]) 193 e = e + 1 194 } 195 i = i + 1 196 } 197 mb = mb + 1 198 } 199 var w: i64 = 0 200 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 } 201 w = 0 202 while w < E*V { Wo[w] = Wo[w] - fxmul(lr, gW[w]/8); w = w + 1 } 203 ep = ep + 1 204 } 205 206 // ---- T1: attention P(correct) on a held-out batch ---- 207 var accsum: i64 = 0 208 var t: i64 = 0 209 while t < 64 { 210 gen_example(900000 + t, keyid, valid, qb, gb) 211 fwd_attn(Kemb, Vemb, Wo, keyid, valid, qb[0], att, ctx, scr, probs, lg, es) 212 accsum = accsum + probs[gb[0]] 213 t = t + 1 214 } 215 let attP: i64 = accsum / 64 216 g_puts(" ATTENTION P(correct)="); g_pn(attP); g_puts(" (ONE="); g_pn(ONE); g_puts(")\n" as *u8) 217 total = total + 1 218 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) } 219 220 // ---- T2: NO-ATTENTION baseline (query-blind mean of the values) ---- 221 var baccsum: i64 = 0 222 t = 0 223 while t < 64 { 224 gen_example(900000 + t, keyid, valid, qb, gb) 225 var e: i64 = 0 226 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 } 227 var v: i64 = 0 228 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 } 229 smax(lg, probs, V, es) 230 baccsum = baccsum + probs[gb[0]] 231 t = t + 1 232 } 233 let baseP: i64 = baccsum / 64 234 g_puts(" NO-ATTENTION baseline P(correct)="); g_pn(baseP); g_puts(" (chance = ONE/V = "); g_pn(ONE/V); g_puts(")\n" as *u8) 235 total = total + 1 236 var t2: i64 = 0 237 if baseP < ONE/2 { if attP > baseP + ONE/8 { t2 = 1 } } 238 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) } 239 240 // ---- T4: deterministic ---- 241 var acc2: i64 = 0 242 t = 0 243 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 } 244 total = total + 1 245 if acc2/64 == attP { pass = pass + 1; ck("T4 deterministic" as *u8, 1) } else { ck("T4 deterministic" as *u8, 0) } 246 247 g_puts("---- nx_reader_attn_gate: passed "); g_pn(pass); g_puts(" / "); g_pn(total); g_puts("\n" as *u8) 248 // MIGRATED onto nx_gate_verdict by nx_gate_dry_apply (D001, minimal form): every check 249 // row above is untouched, so the PASS/FAIL vector cannot change; only the hand-rolled 250 // verdict emission is replaced by the ONE shared base class. Proven by nx_gate_migrate verify. 251 let ctr__dry: *i64 = gv_ctr() 252 ctr__dry[0] = pass 253 ctr__dry[1] = total 254 let rc__dry: i64 = gv_verdict("READER-ATTN-GATE" as *u8, ctr__dry, "dot-product attention learns context-alignment (query->key->value lookup) on the sovereign integer stack; the capability the isolated-candidate reader lacks." as *u8) 255 sys_exit(rc__dry) 256 return rc__dry 257}