nx_reader_span_gate.nx source

1// nx_reader_span_gate.nx -- RUNG 2 of the NEURAL PASSAGE READER arc (deep-research R2): the READER-SHAPED 2// MODEL trains END-TO-END on the verified sovereign no-float stack. Composes nx_nofloat_autograd's PROVEN 3// transformer block (the exact blk_fwd shape from the green blockfloat gates: embed -> rmsnorm -> Q/K/V -> 4// RoPE -> QK^T -> causal softmax -> A.V -> Wo -> residual -> rmsnorm -> SwiGLU FFN -> residual) + a SPAN HEAD 5// (matmul_nt of a trained probe against the contextual states -> [1,L] position logits) + fused softmax-CE. 6// 7// TASK (question-conditioned span pointing = the reader mechanism): passage = 3 (key,value) pairs in random 8// order; question = one key token; GOLD SPAN = the position where the queried key occurs in the passage. To 9// point there the model must ALIGN the question to the context (attention); a QUESTION-BLIND ablation (same 10// params/budget, attention removed -> h_i sees only its own token, exactly the isolated-candidate reader that 11// plateaued at 251/148/101) can at best fire on one fixed key = ~333 permille -> liar-kill. 12// T1 reader (attention block) span accuracy > 700 permille held-out 13// T2 question-blind ablation < 450 permille AND reader beats it by > 250 (liar-kill) 14// T3 descent: one nfa_backward+nfa_sgd step on fresh weights DECREASES that example's loss (real gradients) 15// T4 deterministic re-eval (integer stack -> bit-identical) 16// expect_exit: 0 license_tier: ORIGINAL Sovereign: nx_nofloat_autograd + nx_syscalls. 17import "nx_nofloat_autograd.nx" 18import "nx_syscalls.nx" 19 20func rs_puts(s: *u8) -> i64 { var n: i64=0; while s[n]!=(0 as u8){n=n+1} sys_write(1,s,n); return 0 } 21func rs_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 } 22func rs_ck(name: *u8, c: i64) -> i64 { if c==1 { rs_puts(" PASS " as *u8) } else { rs_puts(" FAIL " as *u8) } rs_puts(name); rs_puts("\n" as *u8); return c } 23 24const RS_V: i64 = 7 // tokens: keys 0..2, values 3..5, SEP 6 25const RS_E: i64 = 8 // model dim (even, for RoPE) 26const RS_F: i64 = 16 // FFN dim 27const RS_L: i64 = 8 // sequence: [q, SEP, k,v, k,v, k,v] 28const RS_SCALE: i64 = 23170 // 1/sqrt(E) in Q16 29// flat weight layout 30const OFF_EMB: i64 = 0 // [V,E] 56 31const OFF_WQ: i64 = 56 // [E,E] 64 32const OFF_WK: i64 = 120 33const OFF_WV: i64 = 184 34const OFF_WO: i64 = 248 35const OFF_WG: i64 = 312 // [E,F] 128 36const OFF_WU: i64 = 440 // [E,F] 128 37const OFF_WD: i64 = 568 // [F,E] 128 38const OFF_U: i64 = 696 // [1,E] 8 (span probe) 39const RS_NP: i64 = 704 40 41func rs_lcg(x: i64) -> i64 { return ((x*1103515245+12345) & 2147483647) } 42 43// permutations of (0,1,2), flat 6x3 44func rs_perms(pm: *i64) -> i64 { 45 pm[0]=0;pm[1]=1;pm[2]=2; pm[3]=0;pm[4]=2;pm[5]=1; pm[6]=1;pm[7]=0;pm[8]=2 46 pm[9]=1;pm[10]=2;pm[11]=0; pm[12]=2;pm[13]=0;pm[14]=1; pm[15]=2;pm[16]=1;pm[17]=0 47 return 0 48} 49 50// example #s -> ids[L], tgt[0]=gold span position (the passage position of the queried key) 51func rs_gen(s: i64, pm: *i64, ids: *i64, tgt: *i64) -> i64 { 52 var r: i64 = rs_lcg(s + 1) 53 let pk: i64 = r % 6 // key-order permutation 54 r = rs_lcg(r) 55 let pv: i64 = r % 6 // value-pairing permutation 56 r = rs_lcg(r) 57 let qs: i64 = r % 3 // queried pair slot 58 ids[0] = pm[pk*3+qs] // question = the queried KEY token 59 ids[1] = 6 // SEP 60 var j: i64 = 0 61 while j < 3 { 62 ids[2+2*j] = pm[pk*3+j] // key token of pair j 63 ids[3+2*j] = 3 + pm[pv*3+j] // value token of pair j 64 j = j + 1 65 } 66 tgt[0] = 2 + 2*qs // gold = position of the queried key in the passage 67 return 0 68} 69 70// deterministic small init (distinct rows; different prime pattern per matrix to break symmetry) 71func rs_init(W: *i64) -> i64 { 72 var i: i64 = 0 73 while i < 56 { W[OFF_EMB+i] = (((i*7+3)%11)-5)*(NFA_Q16/16); i = i + 1 } 74 i = 0 75 while i < 64 { W[OFF_WQ+i] = (((i*13)%9)-4)*(NFA_Q16/24); i = i + 1 } 76 i = 0 77 while i < 64 { W[OFF_WK+i] = (((i*17)%9)-4)*(NFA_Q16/24); i = i + 1 } 78 i = 0 79 while i < 64 { W[OFF_WV+i] = (((i*19)%9)-4)*(NFA_Q16/24); i = i + 1 } 80 i = 0 81 while i < 64 { W[OFF_WO+i] = (((i*23)%9)-4)*(NFA_Q16/24); i = i + 1 } 82 i = 0 83 while i < 128 { W[OFF_WG+i] = (((i*29)%9)-4)*(NFA_Q16/32); i = i + 1 } 84 i = 0 85 while i < 128 { W[OFF_WU+i] = (((i*31)%9)-4)*(NFA_Q16/32); i = i + 1 } 86 i = 0 87 while i < 128 { W[OFF_WD+i] = (((i*37)%9)-4)*(NFA_Q16/32); i = i + 1 } 88 i = 0 89 while i < 8 { W[OFF_U+i] = (((i*5+1)%7)-3)*(NFA_Q16/8); i = i + 1 } 90 return 0 91} 92 93// forward the READER on the tape. useattn=0 -> QUESTION-BLIND ablation (attention sublayer removed; h_i sees 94// only its own token = the isolated-candidate reader). Fills out[0]=loss node, out[1]=span node, out[2..10]= 95// leaf node ids (Emb,Wq,Wk,Wv,Wo,Wg,Wu,Wd,U) for gradient extraction. 96func rs_fwd(tape: *i64, vals: *i64, st: *i64, W: *i64, ids: *i64, tgt: *i64, useattn: i64, out: *i64) -> i64 { 97 st[0] = 0 98 st[1] = 0 99 let nEmb: i64 = nfa_leaf(tape, vals, st, RS_V, RS_E, W, OFF_EMB) 100 let nWq: i64 = nfa_leaf(tape, vals, st, RS_E, RS_E, W, OFF_WQ) 101 let nWk: i64 = nfa_leaf(tape, vals, st, RS_E, RS_E, W, OFF_WK) 102 let nWv: i64 = nfa_leaf(tape, vals, st, RS_E, RS_E, W, OFF_WV) 103 let nWo: i64 = nfa_leaf(tape, vals, st, RS_E, RS_E, W, OFF_WO) 104 let nWg: i64 = nfa_leaf(tape, vals, st, RS_E, RS_F, W, OFF_WG) 105 let nWu: i64 = nfa_leaf(tape, vals, st, RS_E, RS_F, W, OFF_WU) 106 let nWd: i64 = nfa_leaf(tape, vals, st, RS_F, RS_E, W, OFF_WD) 107 let nU: i64 = nfa_leaf(tape, vals, st, 1, RS_E, W, OFF_U) 108 let nX: i64 = nfa_embed(tape, vals, st, nEmb, ids, RS_L) 109 var nH: i64 = nX 110 if useattn == 1 { 111 let nXn: i64 = nfa_rmsnorm_rows(tape, vals, st, nX) 112 let nQ: i64 = nfa_matmul(tape, vals, st, nXn, nWq) 113 let nK: i64 = nfa_matmul(tape, vals, st, nXn, nWk) 114 let nV: i64 = nfa_matmul(tape, vals, st, nXn, nWv) 115 let nQr: i64 = nfa_rope(tape, vals, st, nQ) 116 let nKr: i64 = nfa_rope(tape, vals, st, nK) 117 let nS: i64 = nfa_matmul_nt(tape, vals, st, nQr, nKr) 118 let nSs: i64 = nfa_cmul(tape, vals, st, nS, RS_SCALE) 119 let nA: i64 = nfa_softmax_rows(tape, vals, st, nSs, 1) 120 let nO: i64 = nfa_matmul(tape, vals, st, nA, nV) 121 let nOp: i64 = nfa_matmul(tape, vals, st, nO, nWo) 122 nH = nfa_vadd(tape, vals, st, nX, nOp) 123 } 124 let nHn: i64 = nfa_rmsnorm_rows(tape, vals, st, nH) 125 let nG: i64 = nfa_matmul(tape, vals, st, nHn, nWg) 126 let nU2: i64 = nfa_matmul(tape, vals, st, nHn, nWu) 127 let nSg: i64 = nfa_silu(tape, vals, st, nG) 128 let nHs: i64 = nfa_hadamard(tape, vals, st, nSg, nU2) 129 let nD: i64 = nfa_matmul(tape, vals, st, nHs, nWd) 130 let nY: i64 = nfa_vadd(tape, vals, st, nH, nD) 131 let nYn: i64 = nfa_rmsnorm_rows(tape, vals, st, nY) 132 let nSpan: i64 = nfa_matmul_nt(tape, vals, st, nU, nYn) // [1,E].[L,E]^T -> [1,L] position logits 133 let nLoss: i64 = nfa_softce_rows(tape, vals, st, nSpan, tgt) 134 out[0] = nLoss 135 out[1] = nSpan 136 out[2] = nEmb 137 out[3] = nWq 138 out[4] = nWk 139 out[5] = nWv 140 out[6] = nWo 141 out[7] = nWg 142 out[8] = nWu 143 out[9] = nWd 144 out[10] = nU 145 return nLoss 146} 147 148// accumulate the 9 leaves' grads (tape arena) into the flat G accumulator (weight layout offsets). 149// NO mmap here -- runs inside the training loop (flat ifs per the nx_cc no-else-chain gotcha). 150func rs_grab(tape: *i64, grads: *i64, out: *i64, G: *i64) -> i64 { 151 var k: i64 = 0 152 while k < 9 { 153 var sz: i64 = 64 154 var of: i64 = 0 155 if k == 0 { sz = 56; of = OFF_EMB } 156 if k == 1 { of = OFF_WQ } 157 if k == 2 { of = OFF_WK } 158 if k == 3 { of = OFF_WV } 159 if k == 4 { of = OFF_WO } 160 if k == 5 { sz = 128; of = OFF_WG } 161 if k == 6 { sz = 128; of = OFF_WU } 162 if k == 7 { sz = 128; of = OFF_WD } 163 if k == 8 { sz = 8; of = OFF_U } 164 let nid: i64 = out[2+k] 165 let goff: i64 = tape[7*nid+6] 166 var i: i64 = 0 167 while i < sz { G[of+i] = G[of+i] + grads[goff+i]; i = i + 1 } 168 k = k + 1 169 } 170 return 0 171} 172 173// train a reader (attention or ablated) for EP epochs of B examples; returns nothing (W updated in place) 174func rs_train(tape: *i64, vals: *i64, grads: *i64, st: *i64, W: *i64, pm: *i64, useattn: i64, EP: i64, lr: i64, verbose: i64) -> i64 { 175 let ids: *i64 = sys_mmap(RS_L*8) as *i64 176 let tgt: *i64 = sys_mmap(8) as *i64 177 let out: *i64 = sys_mmap(12*8) as *i64 178 let G: *i64 = sys_mmap(RS_NP*8) as *i64 179 let B: i64 = 8 180 var ep: i64 = 0 181 while ep < EP { 182 var z: i64 = 0 183 while z < RS_NP { G[z] = 0; z = z + 1 } 184 var lsum: i64 = 0 185 var mb: i64 = 0 186 while mb < B { 187 rs_gen(ep*B+mb, pm, ids, tgt) 188 let nLoss: i64 = rs_fwd(tape, vals, st, W, ids, tgt, useattn, out) 189 lsum = lsum + nfa_val(tape, vals, nLoss, 0) 190 nfa_backward(tape, vals, grads, st[0], nLoss) 191 rs_grab(tape, grads, out, G) 192 mb = mb + 1 193 } 194 var w: i64 = 0 195 while w < RS_NP { G[w] = G[w] / B; w = w + 1 } 196 nfa_sgd(W, G, RS_NP, lr) 197 if verbose == 1 { if ep % 500 == 0 { rs_puts(" epoch "); rs_pn(ep); rs_puts(" loss_milli="); rs_pn((lsum/B)*1000/NFA_Q16); rs_puts("\n" as *u8) } } 198 ep = ep + 1 199 } 200 return 0 201} 202 203// held-out span accuracy (permille over 64 examples) of a trained reader 204func rs_eval(tape: *i64, vals: *i64, st: *i64, W: *i64, pm: *i64, useattn: i64) -> i64 { 205 let ids: *i64 = sys_mmap(RS_L*8) as *i64 206 let tgt: *i64 = sys_mmap(8) as *i64 207 let out: *i64 = sys_mmap(12*8) as *i64 208 var hits: i64 = 0 209 var t: i64 = 0 210 while t < 64 { 211 rs_gen(900000 + t, pm, ids, tgt) 212 rs_fwd(tape, vals, st, W, ids, tgt, useattn, out) 213 let nSpan: i64 = out[1] 214 var best: i64 = 0 215 var bv: i64 = nfa_val(tape, vals, nSpan, 0) 216 var j: i64 = 1 217 while j < RS_L { let v: i64 = nfa_val(tape, vals, nSpan, j); if v > bv { bv = v; best = j } j = j + 1 } 218 if best == tgt[0] { hits = hits + 1 } 219 t = t + 1 220 } 221 return hits * 1000 / 64 222} 223 224func main() -> i64 { 225 rs_puts("nx_reader_span_gate (RUNG 2: the READER MODEL -- transformer block + span head -- trains end-to-end on the verified nfa stack)\n" as *u8) 226 var pass: i64 = 0 227 var total: i64 = 0 228 let tape: *i64 = sys_mmap(4096*7*8) as *i64 229 let vals: *i64 = sys_mmap(131072*8) as *i64 230 let grads: *i64 = sys_mmap(131072*8) as *i64 231 let st: *i64 = sys_mmap(2*8) as *i64 232 let pm: *i64 = sys_mmap(18*8) as *i64 233 rs_perms(pm) 234 let EP: i64 = 2500 235 let lr: i64 = 4096 236 237 // ---- T3 FIRST (on fresh weights; untrained -> non-trivial): one backward+sgd step must cut the loss ---- 238 let W3: *i64 = sys_mmap(RS_NP*8) as *i64 239 rs_init(W3) 240 let ids3: *i64 = sys_mmap(RS_L*8) as *i64 241 let tgt3: *i64 = sys_mmap(8) as *i64 242 let out3: *i64 = sys_mmap(12*8) as *i64 243 let G3: *i64 = sys_mmap(RS_NP*8) as *i64 244 var z3: i64 = 0 245 while z3 < RS_NP { G3[z3] = 0; z3 = z3 + 1 } 246 rs_gen(777, pm, ids3, tgt3) 247 let nL0: i64 = rs_fwd(tape, vals, st, W3, ids3, tgt3, 1, out3) 248 let loss0: i64 = nfa_val(tape, vals, nL0, 0) 249 nfa_backward(tape, vals, grads, st[0], nL0) 250 rs_grab(tape, grads, out3, G3) 251 nfa_sgd(W3, G3, RS_NP, lr) 252 let nL1: i64 = rs_fwd(tape, vals, st, W3, ids3, tgt3, 1, out3) 253 let loss1: i64 = nfa_val(tape, vals, nL1, 0) 254 rs_puts(" DESCENT loss_milli "); rs_pn(loss0*1000/NFA_Q16); rs_puts(" -> "); rs_pn(loss1*1000/NFA_Q16); rs_puts(" (one nfa_backward+nfa_sgd step; must fall)\n" as *u8) 255 total = total + 1 256 var t3ok: i64 = 0 257 if loss1 < loss0 { if loss0 > NFA_Q16/16 { t3ok = 1 } } 258 if t3ok == 1 { pass = pass + 1; rs_ck("T3 real gradients: one step cuts the loss" as *u8, 1) } else { rs_ck("T3 real gradients: one step cuts the loss" as *u8, 0) } 259 260 // ---- train the READER (attention block) ---- 261 rs_puts(" training READER (attention block + span head, batch-GD)...\n" as *u8) 262 let W: *i64 = sys_mmap(RS_NP*8) as *i64 263 rs_init(W) 264 rs_train(tape, vals, grads, st, W, pm, 1, EP, lr, 1) 265 let accA: i64 = rs_eval(tape, vals, st, W, pm, 1) 266 rs_puts(" READER (attention) held-out span accuracy = "); rs_pn(accA); rs_puts(" permille\n" as *u8) 267 total = total + 1 268 if accA > 700 { pass = pass + 1; rs_ck("T1 reader points at the answer span (>700)" as *u8, 1) } else { rs_ck("T1 reader points at the answer span (>700)" as *u8, 0) } 269 270 // ---- train the QUESTION-BLIND ablation (identical budget/init; attention removed) ---- 271 rs_puts(" training ABLATION (question-blind: attention removed, same params/budget)...\n" as *u8) 272 let W2: *i64 = sys_mmap(RS_NP*8) as *i64 273 rs_init(W2) 274 rs_train(tape, vals, grads, st, W2, pm, 0, EP, lr, 0) 275 let accB: i64 = rs_eval(tape, vals, st, W2, pm, 0) 276 rs_puts(" ABLATION (question-blind) held-out span accuracy = "); rs_pn(accB); rs_puts(" permille (informed ceiling ~333)\n" as *u8) 277 total = total + 1 278 var t2ok: i64 = 0 279 if accB < 450 { if accA > accB + 250 { t2ok = 1 } } 280 if t2ok == 1 { pass = pass + 1; rs_ck("T2 question-blind ablation fails / reader wins (liar-kill)" as *u8, 1) } else { rs_ck("T2 question-blind ablation fails / reader wins (liar-kill)" as *u8, 0) } 281 282 // ---- T4 deterministic re-eval ---- 283 let accA2: i64 = rs_eval(tape, vals, st, W, pm, 1) 284 total = total + 1 285 if accA2 == accA { pass = pass + 1; rs_ck("T4 deterministic re-eval (bit-identical)" as *u8, 1) } else { rs_ck("T4 deterministic re-eval (bit-identical)" as *u8, 0) } 286 287 rs_puts("---- nx_reader_span_gate: passed "); rs_pn(pass); rs_puts(" / "); rs_pn(total); rs_puts("\n" as *u8) 288 if pass == total { rs_puts("READER RUNG2 GREEN -- the reader-shaped model (transformer block + span head + softmax-CE) trains END-TO-END on the sovereign no-float autograd and does QUESTION-CONDITIONED span pointing; the question-blind (isolated-candidate) reader provably cannot.\n" as *u8); return 0 } 289 rs_puts("RED -- reader rung 2 not fully passed (see numbers)\n" as *u8) 290 return 1 291}